ending mode: exact vs natural with auto outro

app.py

@@ -1,6 +1,7 @@
 import json
 import os
 import random
+import re
 import secrets
 import subprocess
 from datetime import datetime, timezone
@@ -140,7 +141,7 @@ def _upload_to_community(wav_path, meta_dict, song_id):
 
 
 @spaces.GPU(duration=120)
-def generate(style, lyrics, cot, seed, cfg_scale, duration):
+def generate(style, lyrics, cot, seed, cfg_scale, duration, ending):
     pipeline = get_pipeline()
     if seed is None or seed < 1:
         try:
@@ -153,13 +154,23 @@ def generate(style, lyrics, cot, seed, cfg_scale, duration):
     except Exception:
         duration = 120
     duration = max(30, min(duration, 360))
+    exact = str(ending or "Exact").lower().startswith("exact")
+    if not exact:
+        # Natural: let the model finish on its own; close with a real [Outro]
+        # tag (official YuE2 vocabulary) so it composes a proper ending.
+        if not re.search(r"\[(outro|end|ending)\]", lyrics or "", re.I):
+            lyrics = (lyrics or "").rstrip() + "\n[Outro]"
     gen_kwargs = dict(style=style, lyrics=lyrics, cot=cot, seed=seed)
     if cfg_scale and float(cfg_scale) != 1.0:
         gen_kwargs["cfg_scale"] = float(cfg_scale)
-    # exact duration: force a hair MORE than needed (~25 tokens/sec),
-    # then sample-trim to the number so players display it right
-    n_tokens = duration * 25 + 25
-    gen_kwargs["semantic_sampling"] = {"min_tokens": n_tokens, "max_tokens": n_tokens}
+    if exact:
+        # exact duration: force a hair MORE than needed (~25 tokens/sec),
+        # then sample-trim to the number so players display it right
+        n_tokens = duration * 25 + 25
+        gen_kwargs["semantic_sampling"] = {"min_tokens": n_tokens, "max_tokens": n_tokens}
+    else:
+        # natural duration: only cap the maximum, model ends when it ends
+        gen_kwargs["semantic_sampling"] = {"max_tokens": duration * 25}
     song = pipeline(**gen_kwargs)
 
     song_id = secrets.token_hex(6)
@@ -167,7 +178,8 @@ def generate(style, lyrics, cot, seed, cfg_scale, duration):
     mp3_path = os.path.join(OUT, f"song_{song_id}.mp3")
     wav_path = os.path.join(OUT, f"{song_id}.wav")
     song.save(flac_path)
-    _finish_length(flac_path, duration)
+    if exact:
+        _finish_length(flac_path, duration)
     flac_to_mp3(flac_path, mp3_path)
     flac_to_wav(flac_path, wav_path)
     ogg_path = os.path.join(OUT, f"{song_id}.ogg")
@@ -195,7 +207,7 @@ def generate(style, lyrics, cot, seed, cfg_scale, duration):
         "created_at": datetime.now(timezone.utc).isoformat(),
     }, song_id)
 
-    return flac_path, flac_path, mp3_path, ogg_path, f"seed={seed} | {actual:.0f}s (asked {duration}s) | {os.path.getsize(mp3_path) // 1024}KB MP3 320kbps"
+    return flac_path, flac_path, mp3_path, ogg_path, f"seed={seed} | {actual:.0f}s" + (f" (asked {duration}s)" if exact else f" (max {duration}s)") + f" | {os.path.getsize(mp3_path) // 1024}KB MP3 320kbps"
 
 
 with gr.Blocks(title="YuE2-3B Music") as demo:
@@ -209,7 +221,8 @@ with gr.Blocks(title="YuE2-3B Music") as demo:
                 seed = gr.Number(label="Seed", value=0, precision=0)
                 cfg = gr.Number(label="CFG scale (1.0 = default)", value=1.0, precision=2)
             with gr.Row():
-                duration = gr.Slider(30, 360, value=60, step=10, label="Duration (seconds, exact)", scale=4)
+                duration = gr.Slider(30, 360, value=60, step=10, label="Duration (seconds)", scale=3)
+                ending = gr.Radio(["Exact", "Natural"], value="Exact", label="Ending", scale=2)
                 duration_txt = gr.Textbox(value="1:00 min", label="≈ time", scale=1, interactive=False)
             btn = gr.Button("Generate song", variant="primary")
             gr.HTML("<div style='text-align:center;margin:-12px 0 -8px;font-size:0.9rem;'>Don't forget to<a style='margin-left:5px;padding:0;' href='https://huggingface.co/m-a-p/YuE2-3B' target='_blank'>like the model ❤️</a></div>")
@@ -219,7 +232,7 @@ with gr.Blocks(title="YuE2-3B Music") as demo:
             mp3_file = gr.File(label="Download MP3 320kbps")
             ogg_file = gr.File(label="Download OGG (Opus)")
             info = gr.Textbox(label="Info")
-    btn.click(fn=generate, inputs=[style, lyrics, cot, seed, cfg, duration], outputs=[audio, flac_file, mp3_file, ogg_file, info], queue=True)
+    btn.click(fn=generate, inputs=[style, lyrics, cot, seed, cfg, duration, ending], outputs=[audio, flac_file, mp3_file, ogg_file, info], queue=True)
     duration.change(fn=_dur_label, inputs=duration, outputs=duration_txt)
     gr.Markdown(
         "---\nPowered by [YuE2](https://huggingface.co/m-a-p/YuE2-3B) · "