text prompt support

app.py

@@ -30,12 +30,18 @@ device=devicetorch.get(torch)
 cfm, tokenizer, muq, vae = prepare_model(device)
 #cfm = torch.compile(cfm)
 
+def clear_audio():
+    return gr.update(value=None)  # Clears the audio field
+
+def clear_text():
+    return gr.update(value="")  # Clears the text field
+
 #@spaces.GPU
-def infer_music(lrc, ref_audio_path, steps, file_type, max_frames=2048):
+def infer_music(lrc, ref_audio_path, steps, file_type, prompt=None, max_frames=2048):
 
     sway_sampling_coef = -1 if steps < 32 else None
     lrc_prompt, start_time = get_lrc_token(lrc, tokenizer, device)
-    style_prompt = get_style_prompt(muq, ref_audio_path)
+    style_prompt = get_style_prompt(muq, ref_audio_path, prompt)
     negative_style_prompt = get_negative_style_prompt(device)
     latent_prompt = get_reference_latent(device, max_frames)
     print(">0")
@@ -191,7 +197,15 @@ with gr.Blocks(css=css) as demo:
                         elem_classes="lyrics-scroll-box",
                         value="""[00:10.00]Moonlight spills through broken blinds\n[00:13.20]Your shadow dances on the dashboard shrine\n[00:16.85]Neon ghosts in gasoline rain\n[00:20.40]I hear your laughter down the midnight train\n[00:24.15]Static whispers through frayed wires\n[00:27.65]Guitar strings hum our cathedral choirs\n[00:31.30]Flicker screens show reruns of June\n[00:34.90]I'm drowning in this mercury lagoon\n[00:38.55]Electric veins pulse through concrete skies\n[00:42.10]Your name echoes in the hollow where my heartbeat lies\n[00:45.75]We're satellites trapped in parallel light\n[00:49.25]Burning through the atmosphere of endless night\n[01:00.00]Dusty vinyl spins reverse\n[01:03.45]Our polaroid timeline bleeds through the verse\n[01:07.10]Telescope aimed at dead stars\n[01:10.65]Still tracing constellations through prison bars\n[01:14.30]Electric veins pulse through concrete skies\n[01:17.85]Your name echoes in the hollow where my heartbeat lies\n[01:21.50]We're satellites trapped in parallel light\n[01:25.05]Burning through the atmosphere of endless night\n[02:10.00]Clockwork gears grind moonbeams to rust\n[02:13.50]Our fingerprint smudged by interstellar dust\n[02:17.15]Velvet thunder rolls through my veins\n[02:20.70]Chasing phantom trains through solar plane\n[02:24.35]Electric veins pulse through concrete skies\n[02:27.90]Your name echoes in the hollow where my heartbeat lies"""    
                     )
-                    audio_prompt = gr.Audio(label="Audio Prompt", type="filepath", value="./src/prompt/default.wav")
+                    with gr.Group():
+                        gr.HTML("<h5>Generate a song from</h5>")
+                        with gr.Row():
+                            with gr.Tab("Audio"):
+                                audio_prompt = gr.Audio(label="Audio Prompt", type="filepath", value="./src/prompt/default.wav")
+                            with gr.Tab("Text Description"):
+                                text_prompt = gr.Textbox(label="Text Prompt", placeholder="Describe the song")
+                    text_prompt.input(clear_audio, inputs=[], outputs=audio_prompt)
+                    audio_prompt.input(clear_text, inputs=[], outputs=text_prompt)
                     
                 with gr.Column():
                     
@@ -333,7 +347,7 @@ with gr.Blocks(css=css) as demo:
     
     lyrics_btn.click(
         fn=infer_music,
-        inputs=[lrc, audio_prompt, steps, file_type],
+        inputs=[lrc, audio_prompt, steps, file_type, text_prompt],
         outputs=audio_output
     )
 

diffrhythm/infer/infer_utils.py

@@ -51,8 +51,12 @@ def get_negative_style_prompt(device):
     
     return vocal_stlye
 
-def get_style_prompt(model, wav_path):
+@torch.no_grad()
+def get_style_prompt(model, wav_path, prompt):
     mulan = model
+
+    if prompt is not None:
+        return mulan(texts=prompt).half()
     
     ext = os.path.splitext(wav_path)[-1].lower()
     if ext == '.mp3':