Update app.py

app.py

@@ -1,13 +1,9 @@
-"""
-HeartMuLa Gradio App - Music Generation with Lyrics and Tags
-A self-contained Gradio app for Hugging Face Spaces
-"""
-
 import os
 import tempfile
 import torch
 import gradio as gr
 from huggingface_hub import hf_hub_download, snapshot_download
+import spaces
 
 # Download models from HuggingFace Hub on startup
 def download_models():
@@ -44,40 +40,29 @@ def download_models():
     print("All models downloaded successfully!")
     return model_dir
 
+from heartlib import HeartMuLaGenPipeline
 
-# Global pipeline instance
-pipeline = None
-
-
-def load_pipeline():
-    """Load the HeartMuLa pipeline."""
-    global pipeline
-    if pipeline is not None:
-        return pipeline
+model_dir = download_models()
 
-    from heartlib import HeartMuLaGenPipeline
+# Determine device and dtype
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+else:
+    device = torch.device("cpu")
+    dtype = torch.float32
 
-    model_dir = download_models()
-
-    # Determine device and dtype
-    if torch.cuda.is_available():
-        device = torch.device("cuda")
-        dtype = torch.bfloat16
-    else:
-        device = torch.device("cpu")
-        dtype = torch.float32
-
-    print(f"Loading pipeline on {device} with {dtype}...")
-    pipeline = HeartMuLaGenPipeline.from_pretrained(
-        model_dir,
-        device=device,
-        dtype=dtype,
-        version="3B",
-    )
-    print("Pipeline loaded successfully!")
-    return pipeline
+print(f"Loading pipeline on {device} with {dtype}...")
+pipeline = HeartMuLaGenPipeline.from_pretrained(
+    model_dir,
+    device=device,
+    dtype=dtype,
+    version="3B",
+)
+print("Pipeline loaded successfully!")
 
 
+@spaces.GPU
 def generate_music(
     lyrics: str,
     tags: str,
@@ -286,10 +271,5 @@ with gr.Blocks(
     )
 
 
-if __name__ == "__main__":
-    # Preload models on startup
-    print("Initializing HeartMuLa...")
-    load_pipeline()
 
-    # Launch the app
-    demo.launch()
+demo.launch()