Fix: lazy load model inside GPU context for ZeroGPU

app.py

@@ -17,18 +17,20 @@ from sam_audio import SAMAudio, SAMAudioProcessor
 
 # Configuration
 MODEL_NAME = "facebook/sam-audio-small"
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print(f"Loading {MODEL_NAME} on {device}...")
+print(f"Loading {MODEL_NAME} processor...")
 
-# Load Model and Processor
-try:
-    model = SAMAudio.from_pretrained(MODEL_NAME).to(device).eval()
-    processor = SAMAudioProcessor.from_pretrained(MODEL_NAME)
-    print("Model loaded successfully.")
-except Exception as e:
-    print(f"Error loading model. Did you set HF_TOKEN in secrets? Error: {e}")
-    raise e
+# Load Processor only (model will be loaded on GPU when needed)
+processor = SAMAudioProcessor.from_pretrained(MODEL_NAME)
+model = None  # Will be loaded lazily
+
+def get_model():
+    global model
+    if model is None:
+        print(f"Loading model to CUDA...")
+        model = SAMAudio.from_pretrained(MODEL_NAME).to("cuda").eval()
+        print("Model loaded successfully.")
+    return model
 
 def save_audio(tensor, sample_rate):
     """Helper to save torch tensor to a temp file for Gradio output."""
@@ -45,19 +47,22 @@ def separate_audio(audio_path, text_prompt):
     if not audio_path:
         return None, None
     
+    # Load model inside GPU context
+    model = get_model()
+    
     # Process Inputs
     inputs = processor(
         audios=[audio_path], 
         descriptions=[text_prompt]
-    ).to(device)
+    ).to("cuda")
 
     # Inference
     with torch.no_grad():
         result = model.separate(inputs)
 
     # Extract Outputs
-    target_audio = result.target[0]   # The sound you asked for
-    residual_audio = result.residual[0] # Everything else
+    target_audio = result.target[0]
+    residual_audio = result.residual[0]
     
     # Get sampling rate from the processor config
     sr = processor.feature_extractor.sampling_rate
@@ -84,7 +89,7 @@ with gr.Blocks(title="SAM-Audio Demo") as demo:
             input_audio = gr.Audio(label="Upload Input Audio", type="filepath")
             text_prompt = gr.Textbox(
                 label="Text Prompt", 
-                placeholder="e.g., 'dog barking', 'man speaking', 'typing keyboard'",
+                placeholder="e.g., 'drums', 'vocals', 'speech', 'piano'",
                 info="Describe the sound you want to isolate."
             )
             run_btn = gr.Button("Separate Audio", variant="primary")