Update app.py

app.py

@@ -1,90 +1,76 @@
-import torch
 import gradio as gr
+import torch
 import random
-import traceback
+import numpy as np
+
+from diffrhythm.infer.infer_utils import (
+    prepare_model,
+    get_lrc_token,
+    get_audio_style_prompt,
+    get_text_style_prompt,
+    get_negative_style_prompt,
+    get_reference_latent
+)
+
+from diffrhythm.infer.infer import inference
 
-# -------------------------------------------------
-# FORCE CPU (FREE TIER SAFE)
-# -------------------------------------------------
+
+# FORCE CPU FOR FREE TIER
 device = "cpu"
 print("Running in CPU mode (Free tier safe)")
 
-# -------------------------------------------------
-# IMPORT YOUR ORIGINAL UTILITIES
-# -------------------------------------------------
-from diffrhythm.infer.infer_utils import prepare_model, infer_music
-
-# -------------------------------------------------
-# LOAD MODEL (REDUCED MEMORY MODE)
-# -------------------------------------------------
-try:
-    cfm, tokenizer, muq, vae, eval_model, eval_muq = prepare_model(
-        max_frames=2048,   # Reduced from 6144
-        device=device
+MAX_SEED = np.iinfo(np.int32).max
+
+cfm, tokenizer, muq, vae, eval_model, eval_muq = prepare_model(
+    max_frames=2048,
+    device=device
+)
+
+
+def generate_music(lrc_text):
+    torch.manual_seed(0)
+
+    lrc_prompt, start_time, end_frame, song_duration = get_lrc_token(
+        2048, lrc_text, tokenizer, 95, device
+    )
+
+    style_prompt = get_text_style_prompt(muq, "emotional piano")
+    negative_style_prompt = get_negative_style_prompt(device)
+
+    latent_prompt, pred_frames = get_reference_latent(
+        device, 2048, False, None, None, vae
     )
-    print("Model loaded successfully.")
-except Exception as e:
-    print("Model loading failed:")
-    print(traceback.format_exc())
-    raise e
-
-# -------------------------------------------------
-# SAFE GENERATION FUNCTION
-# -------------------------------------------------
-def generate_music(
-    prompt,
-    duration,
-    guidance_scale,
-    seed
-):
-    try:
-        if seed == -1:
-            seed = random.randint(0, 999999)
-
-        torch.manual_seed(int(seed))
-
-        audio = infer_music(
-            cfm=cfm,
-            tokenizer=tokenizer,
-            muq=muq,
-            vae=vae,
-            eval_model=eval_model,
-            eval_muq=eval_muq,
-            prompt=prompt,
-            duration=int(duration),
-            guidance_scale=float(guidance_scale),
-            device=device
-        )
-
-        return audio
-
-    except Exception as e:
-        return f"Error during generation:\n{traceback.format_exc()}"
-
-# -------------------------------------------------
-# GRADIO UI
-# -------------------------------------------------
-with gr.Blocks() as demo:
-
-    gr.Markdown("# 🎵 DiffRhythm (CPU Safe Version)")
-    gr.Markdown("Running on Free CPU mode. Generation will be slow.")
-
-    with gr.Row():
-        prompt = gr.Textbox(label="Music Prompt", placeholder="Describe the music...")
-    
-    with gr.Row():
-        duration = gr.Slider(5, 30, value=10, step=1, label="Duration (seconds)")
-        guidance = gr.Slider(1, 10, value=5, step=0.5, label="Guidance Scale")
-        seed = gr.Number(value=-1, label="Seed (-1 = random)")
-
-    generate_btn = gr.Button("Generate Music")
-
-    output_audio = gr.Audio(label="Generated Audio")
-
-    generate_btn.click(
-        fn=generate_music,
-        inputs=[prompt, duration, guidance, seed],
-        outputs=output_audio
+
+    song = inference(
+        cfm_model=cfm,
+        vae_model=vae,
+        eval_model=eval_model,
+        eval_muq=eval_muq,
+        cond=latent_prompt,
+        text=lrc_prompt,
+        duration=end_frame,
+        style_prompt=style_prompt,
+        negative_style_prompt=negative_style_prompt,
+        steps=10,
+        cfg_strength=3.0,
+        sway_sampling_coef=None,
+        start_time=start_time,
+        file_type="mp3",
+        vocal_flag=False,
+        odeint_method="euler",
+        pred_frames=pred_frames,
+        batch_infer_num=1,
+        song_duration=song_duration
     )
 
-demo.launch()
+    return song
+
+
+demo = gr.Interface(
+    fn=generate_music,
+    inputs=gr.Textbox(lines=10, label="LRC Lyrics"),
+    outputs=gr.Audio(type="filepath")
+)
+
+if __name__ == "__main__":
+    demo.launch()