app.py

"""
NotaGen Gradio Inference App for HuggingFace Spaces

This app provides a simple interface for generating symbolic music using NotaGen.
It accepts three parameters (period, composer, instrumentation) and returns ABC notation.

The ABC notation can then be processed by WeaveMuse locally for XML/PDF conversion.

IMPORTANT: Zero GPU Strategy
----------------------------
- Model initialization and weight downloading happen OUTSIDE @spaces.GPU decorated functions
- Only the actual inference happens inside the GPU-allocated function
- This ensures efficient GPU usage (only during inference, not during setup)
"""

import gradio as gr
import spaces
import torch
import logging
import traceback
import os
from smolagents import Tool
from typing import Optional
from weavemuse.models.notagen.inference import inference_patch, download_model_weights, model


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ============================================================================
# INITIALIZATION PHASE (Outside GPU allocation)
# ============================================================================
# Download model weights and prepare everything BEFORE GPU functions are called
# This ensures GPU is only used for actual inference, not for setup

device = "cuda"
logger.info(f"Preparing NotaGen tool on device: {device}")


def load_model_weights(model_id=None):
    """Load model weights with intelligent quantization support."""
    global model
            
    # Fall back to original weights
    try:
        logger.info("Loading original full-precision model...")
        original_path = download_model_weights(repo_id="ElectricAlexis/NotaGen")
        
        checkpoint = torch.load(original_path, map_location=device, weights_only=False)
        if 'model' in checkpoint:
            state_dict = checkpoint['model']
        else:
            state_dict = checkpoint
            
        model.load_state_dict(state_dict)        
        model.eval()
        logger.info("✅ Original model loaded successfully!")
        return True
        
    except Exception as e:
        logger.error(f"Failed to load model weights: {e}")
        import traceback
        logger.error(traceback.format_exc())
        raise



class SimpleNotaGenTool(Tool):
    """
    Simple tool for symbolic music generation using NotaGen model.
    
    This tool can:
    - Generate ABC notation from text prompts
    - Create music in specific styles and periods
    - Generate compositions for specified instrumentation
    - Handle conditional generation with period-composer-instrumentation prompts
    - Convert ABC to XML, PDF, MIDI, and MP3 formats
    - Generate PDF images for visual display
    
    Note: This is a simplified version without VRAM management.
    """
    
    # Class attributes required by smolagents
    name = "notagen"
    description = (
        "Generates symbolic music in ABC notation format with full conversion capabilities. "
        "Can create compositions only accepts three parameters: musical period, composer, and instrumentation (Use Piano for better results). "
        "composers, and instrumentation. Supports conditional generation with format: "
        "'Period-Composer-Instrumentation' (e.g., 'Romantic-Chopin-Piano'). "
        "Automatically converts to various formats including PDF for visual display."        
    )
    inputs = {
        "period": {
            "type": "string", 
            "description": "Musical period (e.g., Baroque, Classical, Romantic)",
        },
        "composer": {
            "type": "string",
            "description": "Composer style to emulate (e.g., Bach, Mozart, Chopin)",
        },
        "instrumentation": {
            "type": "string",
            "description": "Instruments to use (e.g., Piano, Violin, Orchestra)",
        }
    }
    output_type = "string"
    
    def __init__(
        self, 
        device: str = "auto",
        model_id: str = "manoskary/NotaGenX-Quantized", 
        output_dir: Optional[str] = None,
        **kwargs
    ):
        """
        Initialize NotaGen tool.
        
        Args:
            device: Device to run on ("auto", "cuda", "cpu")
            model_id: NotaGen model ID
            output_dir: Directory for output files
            **kwargs: Additional arguments
        """
        
        # NotaGen is smaller model, estimate VRAM usage
        estimated_vram = 2000.0
        
        super().__init__()
        
        self.output_dir = output_dir or "/tmp/notagen_output"
        self.model_id = model_id
        self.device = device
        
        # Download model weights during initialization (outside GPU function)
        self.download_model_weights(repo_id=model_id)

        os.makedirs(self.output_dir, exist_ok=True)        
        logger.info(f"Simple NotaGen tool initialized")

    def forward(self, period: str, composer: str, instrumentation: str) -> str:
        """
        Generate symbolic music using NotaGen.
        
        Args:
            period: Musical period (e.g., Baroque, Classical, Romantic)
            composer: Composer style to emulate (e.g., Bach, Mozart, Chopin)
            instrumentation: Instruments to use (e.g., Piano, Violin, Orchestra)
            
        Returns:
            Path to generated ABC file or error message
        """
        global model
        global device

        logger.info(f"Generating music: {period}-{composer}-{instrumentation}")
        
        # Create prompt for NotaGen
        prompt = f"{period}-{composer}-{instrumentation}"
        model = model.to(device)
        # Use the inference function
        inference_fn = inference_patch
        if inference_fn is None:
            raise ImportError("inference_patch not available")
            
        # Generate ABC notation (placeholder implementation)
        abc_content = inference_fn(period, composer, instrumentation)
        return abc_content
    
    def download_model_weights(self, repo_id="manoskary/NotaGenX"):
        """
        Download model weights from HuggingFace.
        
        Args:
            repo_id: Repository ID on HuggingFace
        """        
        load_model_weights(model_id=repo_id)
        logger.info("✅ NotaGen model weights downloaded successfully!")


# ============================================================================
# TOOL INITIALIZATION (Outside GPU allocation)
# ============================================================================
# Initialize the tool and download model weights BEFORE the GPU function is called
# This ensures:
# 1. Model weights are downloaded once at startup (not during every inference)
# 2. GPU is only allocated for actual inference, not for downloading/setup
# 3. Zero GPU is used efficiently (shorter GPU allocation times)

try:
    notagen_tool = SimpleNotaGenTool(device=device, model_id="ElectricAlexis/NotaGen")
    logger.info("✅ NotaGen tool initialized successfully!")
except Exception as e:
    logger.error(f"❌ Failed to initialize NotaGen tool: {e}")
    logger.error(traceback.format_exc())
    notagen_tool = None


# ============================================================================
# GPU-ALLOCATED INFERENCE FUNCTION
# ============================================================================
# This function is decorated with @spaces.GPU to request GPU only during execution
# Model is already loaded, so GPU time is minimal and efficient

@spaces.GPU(duration=120)
def generate_abc(period: str, composer: str, instrumentation: str) -> str:
    """
    Generate ABC notation using NotaGen.
    
    This function is decorated with @spaces.GPU to allocate GPU only during inference.
    Model weights are already downloaded and prepared outside this function.
    
    Args:
        period: Musical period (e.g., Baroque, Classical, Romantic)
        composer: Composer style (e.g., Bach, Mozart, Chopin)
        instrumentation: Instruments (e.g., Piano, Violin, Orchestra)
    
    Returns:
        ABC notation string or error message
    """
    if notagen_tool is None:
        error_msg = "❌ NotaGen tool not initialized. Please check server logs."
        logger.error(error_msg)
        return error_msg
    
    try:
        logger.info(f"Generating ABC for: {period}-{composer}-{instrumentation}")
        
        # Call the NotaGen tool's forward method
        # Model is already loaded, this just does the inference
        result = notagen_tool.forward(
            period=period,
            composer=composer,
            instrumentation=instrumentation
        )    
        
        logger.info(f"✅ Successfully generated ABC notation ({len(result)} chars)")
        return result
        
    except Exception as e:
        error_msg = f"❌ Error during generation: {str(e)}\n\n{traceback.format_exc()}"
        logger.error(error_msg)
        return error_msg


# Create Gradio interface
with gr.Blocks(title="NotaGen - Symbolic Music Generation") as demo:
    gr.Markdown("""
    # 🎵 NotaGen - Symbolic Music Generation
    
    Generate symbolic music in ABC notation format using NotaGen.
    
    This space returns only the **ABC notation** as text. For conversion to PDF, XML, or MIDI,
    use the WeaveMuse package locally or through its API.
    
    ### Usage
    1. Select a musical **period** (e.g., Baroque, Classical, Romantic)
    2. Choose a **composer** style (e.g., Bach, Mozart, Chopin)
    3. Specify **instrumentation** (Piano recommended for best results)
    4. Click **Generate** to receive ABC notation
    
    The generated ABC can be processed by WeaveMuse for full music score rendering.
    """)
    
    with gr.Row():
        with gr.Column():
            period_input = gr.Textbox(
                label="Musical Period",
                placeholder="e.g., Classical, Romantic, Baroque",
                value="Classical"
            )
            composer_input = gr.Textbox(
                label="Composer Style",
                placeholder="e.g., Mozart, Chopin, Bach",
                value="Mozart"
            )
            instrumentation_input = gr.Textbox(
                label="Instrumentation",
                placeholder="e.g., Piano, Violin, Orchestra",
                value="Piano"
            )
            generate_btn = gr.Button("🎼 Generate ABC Notation", variant="primary")
        
        with gr.Column():
            output_text = gr.Textbox(
                label="Generated ABC Notation",
                placeholder="ABC notation will appear here...",
                lines=20,
                max_lines=30
            )
    
    gr.Examples(
        examples=[
            ["Classical", "Mozart", "Piano"],
            ["Romantic", "Chopin", "Piano"],
            ["Baroque", "Bach", "Piano"],
            ["Classical", "Beethoven", "Piano"],
            ["Romantic", "Liszt", "Piano"],
        ],
        inputs=[period_input, composer_input, instrumentation_input],
        label="Example Prompts"
    )
    
    # Connect the button to the function
    generate_btn.click(
        fn=generate_abc,
        inputs=[period_input, composer_input, instrumentation_input],
        outputs=output_text,
        api_name="infer"  # Important: This creates the /infer endpoint
    )
    
    gr.Markdown("""
    ---
    ### About
    
    **NotaGen** is a symbolic music generation model that creates music in ABC notation format.
    
    - **ABC Notation**: A text-based music notation format that can be converted to PDF, MIDI, XML, etc.
    - **Model**: Uses the quantized NotaGen model for efficient inference
    - **Integration**: Designed to work seamlessly with WeaveMuse for full music generation pipelines
    
    **Note**: This space only generates ABC notation. For complete score rendering (PDF, MP3, etc.),
    use WeaveMuse locally or via its remote tools.
    """)


# Launch the app
if __name__ == "__main__":
    demo.launch()