Update app.py

app.py

@@ -5,6 +5,12 @@ 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
@@ -12,19 +18,145 @@ import spaces
 import torch
 import logging
 import traceback
-from weavemuse.tools.notagen_tool import NotaGenTool
+import os
+from smolagents import Tool
+from typing import Optional
+from weavemuse.models.notagen.inference import inference_patch, download_model_weights
+
 
-# Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize the NotaGen tool
-# This will use GPU if available in the HuggingFace Space
+# ============================================================================
+# 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"Initializing NotaGen tool on device: {device}")
+logger.info(f"Preparing NotaGen tool on device: {device}")
+
+
+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
+        """                                            
+        logger.info(f"Generating music: {period}-{composer}-{instrumentation}")
+        
+        # Create prompt for NotaGen
+        prompt = f"{period}-{composer}-{instrumentation}"
+        
+        # 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
+        """        
+        download_model_weights(repo_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 = NotaGenTool(device=device)
+    notagen_tool = SimpleNotaGenTool(device=device)
     logger.info("✅ NotaGen tool initialized successfully!")
 except Exception as e:
     logger.error(f"❌ Failed to initialize NotaGen tool: {e}")
@@ -32,11 +164,20 @@ except Exception as e:
     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)
@@ -54,31 +195,14 @@ def generate_abc(period: str, composer: str, instrumentation: str) -> str:
         logger.info(f"Generating ABC for: {period}-{composer}-{instrumentation}")
         
         # Call the NotaGen tool's forward method
-        # This will return a formatted string with file paths
+        # Model is already loaded, this just does the inference
         result = notagen_tool.forward(
             period=period,
             composer=composer,
             instrumentation=instrumentation
-        )
+        )    
         
-        # Extract just the ABC content from the result
-        # The result contains file paths and other info, but we only need the ABC
-        if "Error:" in result:
-            logger.error(f"Generation error: {result}")
-            return result
-        
-        # Parse the result to extract ABC content
-        # The result format is: "Successfully generated...\n\nABC notation preview:\n```\n{abc}```"
-        if "ABC notation preview:" in result:
-            abc_start = result.find("```\n") + 4
-            abc_end = result.rfind("\n```")
-            if abc_start > 3 and abc_end > abc_start:
-                abc_content = result[abc_start:abc_end]
-                logger.info(f"✅ Successfully generated ABC notation ({len(abc_content)} chars)")
-                return abc_content
-        
-        # If we can't parse it, return the full result
-        logger.warning("Could not parse ABC from result, returning full result")
+        logger.info(f"✅ Successfully generated ABC notation ({len(result)} chars)")
         return result
         
     except Exception as e:
@@ -170,8 +294,4 @@ with gr.Blocks(title="NotaGen - Symbolic Music Generation") as demo:
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",  # Allow external access
-        server_port=7860,        # Default HF Spaces port
-        share=False              # No need to share, HF handles this
-    )
+    demo.launch()