Add explicit Target Duration slider (0–60s) + gen_duration kwarg

- TTSServer.generate gains gen_duration: float = 0.0 (override estimator)
- Gradio app exposes a 'Target duration (s)' slider (0 = auto)
- inference.py --gen-duration help text expanded
- Auto-rescale and end-of-clip patch already active for long outputs

For music / multi-section scenes set the slider to 20–60 s and the
auto-rescale schedule keeps the output safe at any cfg.

app.py

@@ -107,7 +107,8 @@ EXAMPLES: list[tuple[str, str, str]] = [
 
 
 @spaces.GPU(duration=120)
-def on_generate(prompt: str, audio_ref, cfg: float, stg: float, dur_mult: float, seed: int):
+def on_generate(prompt: str, audio_ref, cfg: float, stg: float, dur_mult: float,
+                gen_dur: float, seed: int):
     if not prompt or not prompt.strip():
         raise gr.Error("Prompt is empty.")
     t0 = time.time()
@@ -119,6 +120,7 @@ def on_generate(prompt: str, audio_ref, cfg: float, stg: float, dur_mult: float,
         voice_ref=ref_path,
         cfg_scale=cfg, stg_scale=stg,
         duration_multiplier=dur_mult, seed=int(seed),
+        gen_duration=float(gen_dur),
     )
     elapsed = time.time() - t0
     logging.info(f"Generated in {elapsed:.2f}s -> {output}")
@@ -159,7 +161,11 @@ with gr.Blocks(
             with gr.Accordion("Inference settings", open=False):
                 cfg_slider = gr.Slider(1.0, 10.0, value=2.5, step=0.5, label="CFG scale")
                 stg_slider = gr.Slider(0.0, 5.0, value=1.5, step=0.5, label="STG scale")
-                dur_slider = gr.Slider(0.8, 2.0, value=1.1, step=0.05, label="Duration ×")
+                dur_slider = gr.Slider(0.8, 2.0, value=1.1, step=0.05,
+                                       label="Duration × (only used when target duration = 0)")
+                gen_dur_slider = gr.Slider(0.0, 60.0, value=0.0, step=1.0,
+                                           label="Target duration (s) — 0 = auto from prompt; "
+                                                 "set higher (≥20s) for long-form music or scenes")
                 seed_input = gr.Number(value=42, label="Seed", precision=0)
             audio_out = gr.Audio(label="Generated audio", type="filepath")
             with gr.Accordion("Prompt writing guide", open=False):
@@ -176,7 +182,8 @@ with gr.Blocks(
 
     gen_btn.click(
         on_generate,
-        inputs=[prompt_box, audio_ref, cfg_slider, stg_slider, dur_slider, seed_input],
+        inputs=[prompt_box, audio_ref, cfg_slider, stg_slider,
+                dur_slider, gen_dur_slider, seed_input],
         outputs=[audio_out],
     )
 
@@ -185,15 +192,16 @@ with gr.Blocks(
     gr.Examples(
         label="🎬 Click any row to generate a sample",
         examples=[
-            [name, prompt, voice_path, 2.5, 1.5, 1.1, 42]
+            [name, prompt, voice_path, 2.5, 1.5, 1.1, 0.0, 42]
             for name, voice_path, prompt in EXAMPLES
         ],
         example_labels=[name for name, _, _ in EXAMPLES],
         inputs=[gr.Textbox(visible=False, label="Scene"),
                 prompt_box, audio_ref,
-                cfg_slider, stg_slider, dur_slider, seed_input],
+                cfg_slider, stg_slider, dur_slider, gen_dur_slider, seed_input],
         outputs=[audio_out],
-        fn=lambda _name, prompt, ref, cfg, stg, dur, seed: on_generate(prompt, ref, cfg, stg, dur, seed),
+        fn=lambda _name, prompt, ref, cfg, stg, dur, gen_dur, seed: on_generate(
+            prompt, ref, cfg, stg, dur, gen_dur, seed),
         cache_examples=False,
         run_on_click=True,
         examples_per_page=20,

src/inference.py

@@ -230,7 +230,10 @@ def parse_args():
     p.add_argument("--output", default="tts_output.wav")
 
     p.add_argument("--ref-duration", type=float, default=10.0, help="Seconds of voice reference to use")
-    p.add_argument("--gen-duration", type=float, default=0.0, help="Target duration (0=auto)")
+    p.add_argument("--gen-duration", type=float, default=0.0,
+                   help="Target output duration in seconds (0 = auto from prompt + multiplier). "
+                        "Set explicitly for long-form prompts (e.g. --gen-duration 30 for music). "
+                        "Outputs >20.5s automatically engage the end-of-clip silence-prior patch.")
     p.add_argument("--pad-start", type=float, default=0.0,
                    help="Prepend N seconds of silent padding, trimmed after decode (use 0 for clean starts)")
     p.add_argument("--speed", type=float, default=1.0)
@@ -260,7 +263,9 @@ def parse_args():
     p.add_argument("--cfg-scale", type=float, default=None, help="CFG scale (auto: 1.0 distilled, 7.0 dev)")
     p.add_argument("--stg-scale", type=float, default=None, help="STG scale (auto: 0.0 distilled, 1.0 dev)")
     p.add_argument("--stg-block", type=int, default=29, help="Block index for STG perturbation")
-    p.add_argument("--rescale-scale", type=float, default=None, help="Rescale (auto: 0.0 distilled, 0.7 dev)")
+    p.add_argument("--rescale-scale", type=float, default=None,
+                   help="Latent CFG std-rescale (default auto: cfg-aware schedule that prevents "
+                        "output clipping at high cfg; pass any float in [0,1] to override).")
     p.add_argument("--modality-scale", type=float, default=None, help="Modality (auto: 1.0 distilled, 3.0 dev)")
     p.add_argument("--cfg-clamp", type=float, default=0.0, help="Clamp guided pred std to N * cond std (0=disabled)")
     p.add_argument("--steps", type=int, default=None, help="Override steps (auto: distilled sigmas / 30 dev)")
@@ -324,7 +329,9 @@ def main():
     if args.stg_scale is None:
         args.stg_scale = 0.0 if is_distilled else 1.0
     if args.rescale_scale is None:
-        args.rescale_scale = 0.0 if is_distilled else 0.7
+        # Auto cfg-aware rescale: imported from inference_server to keep one source of truth.
+        from inference_server import auto_rescale_for_cfg
+        args.rescale_scale = 0.0 if is_distilled else auto_rescale_for_cfg(args.cfg_scale)
     if args.modality_scale is None:
         args.modality_scale = 1.0 if is_distilled else 3.0
     if args.fps is None:

src/inference_server.py

@@ -60,6 +60,34 @@ def estimate_duration(prompt, multiplier=1.1):
     return max(3.0, round(base * multiplier, 1))
 
 
+def auto_rescale_for_cfg(cfg: float) -> float:
+    """CFG-aware std-rescale schedule that prevents output clipping at high cfg.
+
+    The CFG formula `pred = cond + (cfg-1)*(cond - uncond)` makes pred.std()
+    grow roughly linearly with cfg, which the audio VAE+vocoder render as
+    progressively louder waveforms. By cfg≈3 the output starts hard-clipping
+    at 0 dBFS — and clipped information is unrecoverable in post.
+
+    Empirical sweep on the blues prompt with the back-porch-boogie ref
+    (rescale_scale needed for ≥1 dB peak headroom):
+        cfg=2.5 → 0.2 ;  cfg=3 → 0.6 ; cfg=4 → 0.8 ; cfg=5–8 → 0.8 ; cfg=10 → 1.0
+
+    Piecewise-linear fit through those points; returns 0 below cfg=2 (no CFG
+    even applied at cfg=1), plateaus at 0.8 between cfg=4 and cfg=8 to
+    preserve the "extra punch" of high-CFG generations, and ramps to 1.0 by
+    cfg=10.
+    """
+    if cfg <= 2.0:
+        return 0.0
+    if cfg <= 3.0:
+        return 0.6 * (cfg - 2.0)               # 0 → 0.6
+    if cfg <= 4.0:
+        return 0.6 + 0.2 * (cfg - 3.0)         # 0.6 → 0.8
+    if cfg <= 8.0:
+        return 0.8                              # plateau
+    return min(1.0, 0.8 + 0.1 * (cfg - 8.0))   # 0.8 → 1.0 at cfg=10
+
+
 class TTSServer:
     def __init__(self, checkpoint=None, full_checkpoint=None, gemma_root=None,
                  device="cuda", dtype="bf16", compile_model=True, bnb_4bit=True):
@@ -177,12 +205,23 @@ class TTSServer:
 
     @torch.inference_mode()
     def generate(self, prompt, voice_ref=None, cfg_scale=2.5, stg_scale=1.5,
-                 duration_multiplier=1.1, seed=42, ref_duration=10.0):
-        """Generate audio. Returns (waveform_path, duration_seconds)."""
+                 duration_multiplier=1.1, seed=42, ref_duration=10.0,
+                 rescale_scale="auto", gen_duration: float = 0.0):
+        """Generate audio. Returns (waveform_path, duration_seconds).
+
+        rescale_scale: latent-side CFG std-rescale that prevents clipping at
+            high cfg. Set to "auto" (default) for the cfg-aware schedule, a
+            float in [0, 1] for a fixed override, or 0 to disable.
+        gen_duration: explicit target duration in seconds. 0 (default) → auto
+            from prompt + duration_multiplier; >0 overrides everything else.
+        """
         t_total = time.time()
 
-        # Duration + target shape
-        gen_dur = estimate_duration(prompt, duration_multiplier)
+        # Duration + target shape — explicit gen_duration wins over the estimator.
+        if gen_duration and gen_duration > 0:
+            gen_dur = float(gen_duration)
+        else:
+            gen_dur = estimate_duration(prompt, duration_multiplier)
         fps = 25.0
         n_frames = int(round(gen_dur * fps)) + 1
         n_frames = ((n_frames - 1 + 4) // 8) * 8 + 1
@@ -231,10 +270,13 @@ class TTSServer:
         logging.info(f"Prompt: {time.time()-t0:.2f}s")
 
         # Denoiser
+        resc = auto_rescale_for_cfg(cfg_scale) if rescale_scale == "auto" else float(rescale_scale)
+        if rescale_scale == "auto":
+            logging.info(f"Auto rescale_scale = {resc:.2f} for cfg={cfg_scale}")
         guider = MultiModalGuider(
             params=MultiModalGuiderParams(
                 cfg_scale=cfg_scale, stg_scale=stg_scale,
-                stg_blocks=[29], rescale_scale=0.0, modality_scale=1.0,
+                stg_blocks=[29], rescale_scale=resc, modality_scale=1.0,
             ),
             negative_context=a_ctx_neg,
         )