Update app.py

app.py

@@ -1,65 +1,25 @@
-import torch
 import gradio as gr
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import Response
-from transformers import AutoProcessor, MusicgenForConditionalGeneration
 import scipy.io.wavfile
 import io
-import numpy as np
+from generate import generate_music
 
-# --- INITIALIZATION ---
-MODEL_ID = "facebook/musicgen-small"
-device = "cpu"
+app = FastAPI(title="Basyx MusicGen Hub")
 
-print(f"DEBUG: System boot. Loading {MODEL_ID}...")
-processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = MusicgenForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
-model.to(device)
-print("DEBUG: Model loaded successfully.")
-
-def generate_core(prompt, duration):
-    print(f"!!! TRIGGERED !!! Prompt: {prompt} | Duration: {duration}")
-    if not prompt:
-        return None
-    
+@app.post("/generate")
+async def api_generate(prompt: str = Query(...), duration: int = Query(10)):
     try:
-        duration = min(int(duration), 30)
-        inputs = processor(text=[prompt], padding=True, return_tensors="pt").to(device)
-        
-        # 50 tokens/sec. Reducing guidance for CPU speed.
-        max_tokens = int(duration * 50)
+        sr, audio = generate_music(prompt, duration)
+        if audio is None:
+            raise HTTPException(status_code=500, detail="Generation failed")
         
-        with torch.no_grad():
-            audio_values = model.generate(
-                **inputs, 
-                max_new_tokens=max_tokens,
-                do_sample=True,
-                guidance_scale=3.0
-            )
-
-        sampling_rate = model.config.audio_encoder.sampling_rate
-        audio_data = audio_values[0, 0].cpu().numpy()
-        print("DEBUG: Generation successful.")
-        return sampling_rate, audio_data
+        byte_io = io.BytesIO()
+        scipy.io.wavfile.write(byte_io, rate=sr, data=audio)
+        return Response(content=byte_io.getvalue(), media_type="audio/wav")
     except Exception as e:
-        print(f"ERROR: {str(e)}")
-        return None
-
-# --- FASTAPI FOR N8N ---
-app = FastAPI()
-
-@app.post("/generate")
-async def api_generate(prompt: str = Query(...), duration: int = Query(10)):
-    res = generate_core(prompt, duration)
-    if res is None:
-        raise HTTPException(status_code=500, detail="Generation failed")
-    
-    sr, audio = res
-    byte_io = io.BytesIO()
-    scipy.io.wavfile.write(byte_io, rate=sr, data=audio)
-    return Response(content=byte_io.getvalue(), media_type="audio/wav")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🎵 MusicGen Automation Hub")
     with gr.Row():
@@ -70,8 +30,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column():
             a_out = gr.Audio(label="Output")
     
-    run_btn.click(fn=generate_core, inputs=[p_in, d_in], outputs=a_out)
+    run_btn.click(fn=generate_music, inputs=[p_in, d_in], outputs=a_out)
 
-# Enable the queue and mount
 demo.queue()
 app = gr.mount_gradio_app(app, demo, path="/")