Update SDK version and app.py for HF stability

README.md

@@ -4,7 +4,7 @@ emoji: 🎹
 colorFrom: indigo
 colorTo: blue
 sdk: gradio
-sdk_version: 4.0.0
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -6,14 +6,16 @@ Exposes a Gradio UI and a FastAPI endpoint for remote Vercel integration.
 from __future__ import annotations
 
 import os
+import uuid
+from typing import Any
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
 import torch
-from fastapi import FastAPI, BackgroundTasks
+from fastapi import BackgroundTasks, FastAPI
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
-import gradio as gr
-import soundfile as sf
-import numpy as np
-import uuid
 
 from src.text_to_audio import build_pipeline
 
@@ -22,50 +24,21 @@ MODEL_PRESET = os.getenv("MODEL_PRESET", "musicgen-small")
 USE_4BIT = os.getenv("USE_4BIT", "False").lower() == "true"
 
 print(f"Loading {MODEL_PRESET} (4-bit={USE_4BIT})...")
-pipe = build_pipeline(preset=MODEL_PRESET, use_4bit=USE_4BIT)
-
-# FastAPI Setup
-app = FastAPI(title="MusicSampler API")
+# Force device to cuda if available, otherwise cpu
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = build_pipeline(preset=MODEL_PRESET, use_4bit=USE_4BIT, device_map=device)
 
 class GenRequest(BaseModel):
     prompt: str
     duration: float = 5.0
     model: str = MODEL_PRESET
 
-@app.post("/generate")
-async def api_generate(req: GenRequest, background_tasks: BackgroundTasks):
-    """API Endpoint for DAW-INVADER / Vercel integration."""
-    filename = f"gen_{uuid.uuid4()}.wav"
-    output_path = os.path.join("/tmp", filename)
-    
-    # Generate audio
-    # MusicGen supports 'max_new_tokens' via generate_kwargs
-    # 5 seconds ~ 250 tokens for MusicGen small (50 tokens/sec)
-    tokens = int(req.duration * 50)
-    
-    out = pipe.generate(
-        req.prompt, 
-        generate_kwargs={"max_new_tokens": tokens}
-    )
-    
-    single = out if isinstance(out, dict) else out[0]
-    audio = single["audio"]
-    sr = single["sampling_rate"]
-    
-    if hasattr(audio, "numpy"):
-        arr = audio.numpy()
-    else:
-        arr = np.asarray(audio)
-        
-    sf.write(output_path, arr.T if arr.ndim == 2 else arr, sr)
-    
-    # Clean up file after serving
-    background_tasks.add_task(os.remove, output_path)
-    
-    return FileResponse(output_path, media_type="audio/wav", filename=filename)
-
-# Gradio Interface
+# Gradio Interface functions
 def gradio_gen(prompt, duration):
+    if not prompt or not prompt.strip():
+        return None, "Please enter a prompt."
+    
+    # MusicGen: 5 seconds ~ 250 tokens (50 tokens/sec approx)
     tokens = int(duration * 50)
     out, profile = pipe.generate_with_profile(
         prompt, 
@@ -81,6 +54,7 @@ def gradio_gen(prompt, duration):
         arr = np.asarray(audio)
         
     path = f"/tmp/gradio_{uuid.uuid4()}.wav"
+    # Ensure audio is properly formatted for soundfile
     sf.write(path, arr.T if arr.ndim == 2 else arr, sr)
     return path, f"Generated in {profile.get('time_s', 0):.2f}s (RTF: {profile.get('rtf', 0):.2f})"
 
@@ -99,9 +73,38 @@ with gr.Blocks(title="MusicSampler", theme=gr.themes.Monochrome()) as ui:
             
     btn.click(gradio_gen, inputs=[prompt, duration], outputs=[audio_out, stats])
 
-# Mount Gradio into FastAPI
-app = gr.mount_gradio_app(app, ui, path="/")
+# HF Spaces automatically launches the app defined in app_file if it's sdk: gradio
+# To expose a custom API alongside Gradio, we use the internal FastAPI app.
+app = ui.app 
+
+@app.post("/generate")
+async def api_generate(req: GenRequest, background_tasks: BackgroundTasks):
+    """API Endpoint for DAW-INVADER / Vercel integration."""
+    filename = f"gen_{uuid.uuid4()}.wav"
+    output_path = os.path.join("/tmp", filename)
+    
+    tokens = int(req.duration * 50)
+    out = pipe.generate(
+        req.prompt, 
+        generate_kwargs={"max_new_tokens": tokens}
+    )
+    
+    single = out if isinstance(out, dict) else out[0]
+    audio = single["audio"]
+    sr = single["sampling_rate"]
+    
+    if hasattr(audio, "numpy"):
+        arr = audio.numpy()
+    else:
+        arr = np.asarray(audio)
+        
+    sf.write(output_path, arr.T if arr.ndim == 2 else arr, sr)
+    
+    # Clean up file after serving
+    background_tasks.add_task(os.remove, output_path)
+    
+    return FileResponse(output_path, media_type="audio/wav", filename=filename)
 
+# Standard entry point for HF Spaces
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    ui.launch(server_name="0.0.0.0", server_port=7860)