""" 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()