app.py

import os
import sys
import gradio as gr
import torch
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Set backend before importing pyplot
import matplotlib.pyplot as plt
from PIL import Image
from huggingface_hub import hf_hub_download
import pretty_midi
import librosa
import soundfile as sf
from midi2audio import FluidSynth
import spaces
from tqdm import tqdm

# Remove CPU forcing since we'll use ZeroGPU
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
# torch.set_num_threads(4)

from aria.image_encoder import ImageEncoder
from aria.aria import ARIA

print("=" * 60)
print("ARIA - Art to Music Generator")
print("=" * 60)
print("Initializing model downloads...")
sys.stdout.flush()

# Pre-download all model files at startup
MODEL_FILES = {
    "image_encoder": "image_encoder.pt",
    "continuous_concat": ["continuous_concat/model.pt", "continuous_concat/mappings.pt", "continuous_concat/model_config.pt"],
    "continuous_token": ["continuous_token/model.pt", "continuous_token/mappings.pt", "continuous_token/model_config.pt"],
    "discrete_token": ["discrete_token/model.pt", "discrete_token/mappings.pt", "discrete_token/model_config.pt"]
}

# Create cache directory
CACHE_DIR = os.path.join(os.path.dirname(__file__), "model_cache")
os.makedirs(CACHE_DIR, exist_ok=True)
print(f"Cache directory: {CACHE_DIR}")
# ---------------------------------------------------------------------------
# Make the HuggingFace Hub cache point to the same directory so Transformers
# will instantly find the CLIP weights we snapshot-downloaded instead of
# recopying them at runtime.
os.environ["HF_HOME"] = CACHE_DIR          # HF>=0.23
os.environ["HF_HUB_CACHE"] = CACHE_DIR
# ---------------------------------------------------------------------------
sys.stdout.flush()

# Download and cache all files
cached_files = {}
total_files = sum(len(files) if isinstance(files, list) else 1 for files in MODEL_FILES.values())
current_file = 0

for model_type, files in MODEL_FILES.items():
    print(f"\nProcessing {model_type} model files...")
    sys.stdout.flush()
    
    if isinstance(files, str):
        files = [files]
    
    cached_files[model_type] = []
    for file in files:
        current_file += 1
        print(f"[{current_file}/{total_files}] {file}")
        sys.stdout.flush()
        
        try:
            # Check if file already exists in cache
            repo_id = "vincentamato/aria"
            cached_path = os.path.join(CACHE_DIR, repo_id, file)
            
            if os.path.exists(cached_path):
                file_size = os.path.getsize(cached_path) / (1024 * 1024)  # MB
                print(f"   Found cached file ({file_size:.1f} MB)")
                cached_files[model_type].append(cached_path)
            else:
                print(f"   Downloading from HuggingFace Hub...")
                print(f"   Repository: {repo_id}")
                sys.stdout.flush()
                
                # Download with progress
                cached_path = hf_hub_download(
                    repo_id=repo_id,
                    filename=file,
                    cache_dir=CACHE_DIR,
                    # resume_download=True  # Enable resume if connection drops
                )
                
                if os.path.exists(cached_path):
                    file_size = os.path.getsize(cached_path) / (1024 * 1024)  # MB
                    print(f"Download complete ({file_size:.1f} MB)")
                    cached_files[model_type].append(cached_path)
                else:
                    print(f"Download failed - file not found")
                    
        except Exception as e:
            print(f"   Error with file {file}: {str(e)}")
            sys.stdout.flush()

try:
    print("\nPre-caching CLIP backbone (openai/clip-vit-large-patch14-336)…")
    sys.stdout.flush()
    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id="openai/clip-vit-large-patch14-336",
        cache_dir=CACHE_DIR,
        local_dir_use_symlinks=False,  # make sure files are copied, not linked
        resume_download=True,
    )
    print("CLIP checkpoint cached successfully!\n")
except Exception as clip_err:
    print(f"Warning: failed to pre-cache CLIP model – it will download at runtime. ({clip_err})\n")
    sys.stdout.flush()


print("\n" + "=" * 60)
print("Model file preparation complete!")
print("=" * 60)
sys.stdout.flush()

# Check what we actually got
for model_type, paths in cached_files.items():
    print(f"{model_type}: {len(paths)} files ready")
    
print(f"\nStarting Gradio application...")
sys.stdout.flush()

# Global model cache
models = {}

def create_emotion_plot(valence, arousal):
    """Create a valence-arousal plot with the predicted emotion point"""
    # Create figure in a process-safe way
    fig = plt.figure(figsize=(8, 8), dpi=100)
    ax = fig.add_subplot(111)
    
    # Set background color and style
    plt.style.use('default')  # Use default style instead of seaborn
    fig.patch.set_facecolor('#ffffff')
    ax.set_facecolor('#ffffff')
    
    # Create the coordinate system with a light grid
    ax.grid(True, linestyle='--', alpha=0.2)
    ax.axhline(y=0, color='#666666', linestyle='-', alpha=0.3, linewidth=1)
    ax.axvline(x=0, color='#666666', linestyle='-', alpha=0.3, linewidth=1)
    
    # Plot region
    circle = plt.Circle((0, 0), 1, fill=False, color='#666666', alpha=0.3, linewidth=1.5)
    ax.add_artist(circle)
    
    # Add labels with nice fonts
    font = {'family': 'sans-serif', 'weight': 'medium', 'size': 12}
    label_dist = 1.35  # Increased distance for labels
    ax.text(label_dist, 0, 'Positive', ha='left', va='center', **font)
    ax.text(-label_dist, 0, 'Negative', ha='right', va='center', **font)
    ax.text(0, label_dist, 'Excited', ha='center', va='bottom', **font)
    ax.text(0, -label_dist, 'Calm', ha='center', va='top', **font)
    
    # Plot the point with a nice style
    ax.scatter([valence], [arousal], c='#4f46e5', s=150, zorder=5, alpha=0.8)
    
    # Set limits and labels with more padding
    ax.set_xlim(-1.6, 1.6)
    ax.set_ylim(-1.6, 1.6)
    
    # Format ticks
    ax.set_xticks([-1.5, -1.0, -0.5, 0, 0.5, 1.0, 1.5])
    ax.set_yticks([-1.5, -1.0, -0.5, 0, 0.5, 1.0, 1.5])
    ax.tick_params(axis='both', which='major', labelsize=10)
    
    # Add axis labels with padding
    ax.set_xlabel('Valence', **font, labelpad=15)
    ax.set_ylabel('Arousal', **font, labelpad=15)
    
    # Remove spines
    for spine in ax.spines.values():
        spine.set_visible(False)
    
    # Adjust layout with more padding
    plt.tight_layout(pad=1.5)
    
    # Save to a temporary file and return the path
    temp_path = os.path.join(os.path.dirname(__file__), "output", "emotion_plot.png")
    os.makedirs(os.path.dirname(temp_path), exist_ok=True)
    plt.savefig(temp_path, bbox_inches='tight', dpi=100)
    plt.close(fig)  # Close the figure to free memory
    
    return temp_path

def get_model(conditioning_type):
    """Get or initialize model with specified conditioning"""
    if conditioning_type not in models:
        try:
            # Use cached files
            image_model_path = cached_files["image_encoder"][0]
            midi_model_dir = os.path.dirname(cached_files[conditioning_type][0])
            
            # right after we build CACHE_DIR
            os.environ["HF_HOME"] = CACHE_DIR          # HF  >=0.23
            os.environ["HF_HUB_CACHE"] = CACHE_DIR     # backward compatibility

            models[conditioning_type] = ARIA(
                image_model_checkpoint=image_model_path,
                midi_model_dir=midi_model_dir,
                conditioning=conditioning_type
            )
        except Exception as e:
            print(f"Error initializing {conditioning_type} model: {str(e)}")
            return None
    return models[conditioning_type]

def convert_midi_to_wav(midi_path):
    """Convert MIDI file to WAV using FluidSynth"""
    wav_path = midi_path.replace('.mid', '.wav')
    
    # If WAV file already exists and is newer than MIDI file, use cached version
    if os.path.exists(wav_path) and os.path.getmtime(wav_path) > os.path.getmtime(midi_path):
        return wav_path
    try:
        # Search common soundfont directories for any .sf2 or .sf3 files
        import glob
        soundfont_search_dirs = [
            'C:\\soundfonts\\',                          # Windows user soundfonts
            'C:\\Program Files\\FluidSynth\\sf2\\',      # Windows FluidSynth installation
            '/usr/share/sounds/sf2/',                    # Linux system soundfonts
            '/usr/share/soundfonts/',                    # Linux alternative
            '/usr/local/share/fluidsynth/',              # macOS homebrew
            '/System/Library/Audio/Sounds/Banks/',       # macOS system
        ]
        
        soundfont = None
        for search_dir in soundfont_search_dirs:
            if os.path.exists(search_dir):
                # Look for .sf2 and .sf3 files in this directory
                for extension in ['*.sf2', '*.sf3']:
                    matches = glob.glob(os.path.join(search_dir, extension))
                    if matches:
                        soundfont = matches[0]  # Use first soundfont found
                        break
                if soundfont:
                    break
                
        if soundfont is None:
            print(f"No SoundFont found. Audio playback not available.")
            print(f"MIDI file saved: {midi_path}")
            print(f"To enable audio: Install FluidSynth and place a .sf2 file in C:\\soundfonts\\")
            return None
            
        # Convert MIDI to WAV using direct FluidSynth command
        print(f"Converting MIDI to WAV using SoundFont: {soundfont}")
        
        # Use subprocess to call fluidsynth directly with proper arguments
        import subprocess
        cmd = [
            'fluidsynth',
            '-ni',          # No interactive mode
            '-g', '0.5',    # Gain
            '-r', '44100',  # Sample rate  
            '-F', wav_path, # Output WAV file
            soundfont,      # SoundFont file
            midi_path       # Input MIDI file
        ]
        
        print(f"FluidSynth command: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        
        if result.returncode == 0 and os.path.exists(wav_path):
            print(f"WAV file created: {wav_path}")
            return wav_path
        else:
            print(f"FluidSynth failed with return code: {result.returncode}")
            print(f"Error output: {result.stderr}")
            return None
        
    except Exception as e:
        print(f"Error converting MIDI to WAV: {str(e)}")
        print(f"MIDI file still available: {midi_path}")
        return None

@spaces.GPU(duration=120)
def generate_music(image, conditioning_type, gen_len, temperature, top_p, min_instruments):
    """Generate music from input image"""
    
    print("▶ generate_music entered")
    
    model = get_model(conditioning_type)
    if model is None:
        # IMPORTANT: Return a 3-element tuple, not a dictionary
        return (
            None,  # For emotion_chart
            None,  # For midi_output
            f"Error: Failed to initialize {conditioning_type} model. Please check the logs."
        )
    
    try:
        # Create output directory
        output_dir = os.path.join(os.path.dirname(__file__), "output")
        os.makedirs(output_dir, exist_ok=True)

        # Generate music
        valence, arousal, midi_path = model.generate(
            image_path=image,
            out_dir=output_dir,
            gen_len=gen_len,
            temperature=temperature,
            top_k=-1,
            top_p=float(top_p),
            min_instruments=int(min_instruments)
        )
        
        # Create emotion plot first (needed for both success and failure cases)
        plot_path = create_emotion_plot(valence, arousal)
        
        # Convert MIDI to WAV
        wav_path = convert_midi_to_wav(midi_path)
        if wav_path is None:
            # WAV conversion failed, but we still have MIDI
            result_text = f"""
**Model Type:** {conditioning_type}

**Predicted Emotions:**
- Valence: {valence:.3f} (negative → positive)
- Arousal: {arousal:.3f} (calm → excited)

**Generation Parameters:**
- Temperature: {temperature}
- Top-p: {top_p}
- Min Instruments: {min_instruments}

**Audio Playback Unavailable**
Your music has been generated as a MIDI file, but audio conversion failed.

**MIDI File:** `{os.path.basename(midi_path)}`

**To Enable Audio Playback:**
1. Install FluidSynth: `choco install fluidsynth` (or download from GitHub)
2. Download a SoundFont file (e.g., GeneralUser GS)
3. Place it at: `C:\\soundfonts\\generaluser.sf2`

You can still download and play the MIDI file in any MIDI player!
"""
            # Return MIDI file for download instead of WAV
            return (plot_path, midi_path, result_text)
        
        # Build a nice Markdown result string for successful WAV conversion
        result_text = f"""
**Model Type:** {conditioning_type}

**Predicted Emotions:**
- Valence: {valence:.3f} (negative → positive)
- Arousal: {arousal:.3f} (calm → excited)

**Generation Parameters:**
- Temperature: {temperature}
- Top-p: {top_p}
- Min Instruments: {min_instruments}

Your music has been generated! Click the play button above to listen.
"""

        # RETURN AS A TUPLE
        return (plot_path, wav_path, result_text)
    
    except Exception as e:
        return (
            None,
            None,
            f"Error generating music: {str(e)}"
        )

def generate_music_wrapper(image, conditioning_type, gen_len, note_temp, rest_temp, top_p, min_instruments):
    """Wrapper for generate_music that handles separate temperatures"""
    return generate_music(
        image=image,
        conditioning_type=conditioning_type,
        gen_len=gen_len,
        temperature=[float(note_temp), float(rest_temp)],
        top_p=top_p,
        min_instruments=min_instruments
    )

# Create Gradio interface
with gr.Blocks(title="ARIA - Art to Music Generator") as demo:
    gr.Markdown("""
    # ARIA: Artistic Rendering of Images into Audio
    
    Upload an image and ARIA will analyze its emotional content to generate matching music!
    
    ## How it works:
    1. ARIA first analyzes the emotional content of your image along two dimensions:
       - **Valence**: How positive or negative the emotion is (-1 to 1)
       - **Arousal**: How calm or excited the emotion is (-1 to 1)
    2. These emotions are then used to generate music that matches the mood
    """)

    
    with gr.Row():
        with gr.Column(scale=3):
            # Give the model a simple string path instead of an in-memory array
            # so that PIL.Image.open inside ARIA works correctly across Gradio versions
            image_input = gr.Image(
                label="Upload Image",
                type="filepath"  # return string path rather than numpy array/PIL
            )

            # Quick-start guidance so first-time users immediately know what to do
            gr.Markdown(
                "## Quick Start\n"
                "1. **Click an example artwork below** *or* **upload your own image** above.\n"
                "2. (Optional) Open **Advanced Settings** to fine-tune the generation.\n"
                "3. Hit **Generate Music** to inference the model!"
            )

            # Advanced controls are tucked away inside a collapsible panel to keep the UI clean
            with gr.Accordion("Advanced Settings", open=False):
                gr.Markdown("### Generation Settings")
                
                with gr.Row():
                    with gr.Column():
                        conditioning_type = gr.Radio(
                            choices=["continuous_concat", "continuous_token", "discrete_token"],
                            value="continuous_concat",
                            label="Conditioning Type",
                            info="How the emotional information is incorporated into the music generation"
                        )
                    with gr.Column():
                        gen_len = gr.Slider(
                            minimum=256,
                            maximum=4096,
                            value=1024,
                            step=256,
                            label="Generation Length",
                            info="Number of tokens to generate (longer = more music)"
                        )
                
                with gr.Row():
                    with gr.Column():
                        note_temperature = gr.Slider(
                            minimum=0.1,
                            maximum=2.0,
                            value=1.2,
                            step=0.1,
                            label="Note Temperature",
                            info="Controls randomness of note generation"
                        )
                    with gr.Column():
                        rest_temperature = gr.Slider(
                            minimum=0.1,
                            maximum=2.0,
                            value=1.2,
                            step=0.1,
                            label="Rest Temperature",
                            info="Controls randomness of rest/timing generation"
                        )
                
                with gr.Row():
                    with gr.Column():
                        top_p = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.6,
                            step=0.1,
                            label="Top-p Sampling",
                            info="Nucleus sampling threshold - lower = more focused"
                        )
                    with gr.Column():
                        min_instruments = gr.Slider(
                            minimum=1,
                            maximum=5,
                            value=2,
                            step=1,
                            label="Minimum Instruments",
                            info="Minimum number of instruments in the generated music"
                        )
            
            generate_btn = gr.Button("Generate Music", variant="primary", size="lg")
            
            # Add examples
            # Dynamic path resolution for local vs HF Spaces deployment
            examples_dir = "examples" if os.path.exists("examples") else "ARIA/examples"
            gr.Examples(
                examples=[
                    [f"{examples_dir}/happy.jpg", "continuous_concat", 1024, 1.2, 1.2, 0.6, 2],
                    [f"{examples_dir}/sad.jpeg", "continuous_concat", 1024, 1.2, 1.2, 0.6, 2],
                ],
                inputs=[image_input, conditioning_type, gen_len, note_temperature, rest_temperature, top_p, min_instruments],
                label="Example Artworks (click to load)"
            )
        
        with gr.Column(scale=2):
            emotion_chart = gr.Image(
                label="Predicted Emotions"
            )
            midi_output = gr.Audio(
                label="Generated Music"
            )
            results = gr.Markdown()
    
    gr.Markdown("""
    ## About ARIA
    
    ARIA is a deep learning system that generates music from artwork by:
    1. Using a image-emotion model to extract emotional content from images
    2. Generating matching music using an emotion-conditioned music generation model
    
    The emotion-conditioned MIDI generation model is based on the work by Serkan Sulun et al. in their paper 
    ["Symbolic music generation conditioned on continuous-valued emotions"](https://ieeexplore.ieee.org/document/9762257).
    Original implementation: [github.com/serkansulun/midi-emotion](https://github.com/serkansulun/midi-emotion)
    
    ## Conditioning Types
    
    **continuous_concat (Recommended)**  
    Creates a single vector from valence and arousal values, repeats it across the sequence, and concatenates it with every music token embedding. This approach gives the emotion information *global influence* throughout the entire generation process, allowing the transformer to access emotional context at every timestep. Research shows this method achieves the best performance in both note prediction accuracy and emotional coherence.
    
    **continuous_token**  
    Converts each emotion value (valence and arousal) into separate condition vectors with the same length as music token embeddings, then concatenates them in the sequence dimension. The emotion vectors are inserted at the beginning of the input sequence during generation. This treats emotions similarly to music tokens but can lose influence as the sequence grows longer.
    
    **discrete_token**  
    Quantizes continuous emotion values into 5 discrete bins (very low, low, moderate, high, very high) and converts them into control tokens. These tokens are placed before the music tokens in the sequence. While this represents the current state-of-the-art approach in conditional text generation, it suffers from information loss due to binning and can lose emotional context during longer generations when tokens are truncated.
    """)

    def generate_music_wrapper(image, conditioning_type, gen_len, note_temp, rest_temp, top_p, min_instruments):
        """Wrapper for generate_music that handles separate temperatures"""
        return generate_music(
            image=image,
            conditioning_type=conditioning_type,
            gen_len=gen_len,
            temperature=[float(note_temp), float(rest_temp)],
            top_p=top_p,
            min_instruments=min_instruments
        )

    generate_btn.click(
        fn=generate_music_wrapper,
        inputs=[image_input, conditioning_type, gen_len, note_temperature, rest_temperature, top_p, min_instruments],
        outputs=[emotion_chart, midi_output, results]
    )

# Launch app
demo.launch(share=True, show_api=False)