Expose conservative sampling controls

app.py

@@ -32,6 +32,8 @@ COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
 CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
         "roll": "5", "roll_big": "6", "balloon": "7"}
 SUB = 96
+DEFAULT_TEMPERATURE = 0.8
+DEFAULT_TOP_P = 0.9
 
 _MODELS = {}
 
@@ -262,7 +264,7 @@ def render_audio_plan(mel, title, course, plan=None):
 
 
 def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_planner,
-             sampling, progress=gr.Progress()):
+             sampling, temperature, top_p, progress=gr.Progress()):
     if audio is None:
         raise gr.Error("Please upload an audio file.")
 
@@ -311,6 +313,8 @@ def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_pla
         bpm = float(np.atleast_1d(librosa.beat.beat_track(y=wav, sr=SR)[0])[0])
     if grid is None and abs(bpm - round(bpm)) < 0.06:
         bpm = float(round(bpm))
+    temperature = float(np.clip(temperature or DEFAULT_TEMPERATURE, 0.2, 1.2))
+    top_p = float(np.clip(top_p or DEFAULT_TOP_P, 0.5, 1.0))
 
     plan = None
     P(0.28, "Planning song structure…")
@@ -329,7 +333,7 @@ def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_pla
         g = generate_song_slot(
             M["slot"], mel, grid, course, level=int(level),
             density_bucket=COURSE_DENS[course], greedy=not sampling, seed=0,
-            device=DEVICE, plan=plan,
+            temperature=temperature, top_p=top_p, device=DEVICE, plan=plan,
             on_progress=lambda d, t: P(0.32 + 0.58 * d / max(t, 1),
                                        f"Generating notes (slot-exact)… {d}/{t} windows"))
         tja = write_tja_slots(g, grid, title, course, int(level), wave_name,
@@ -339,7 +343,7 @@ def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_pla
         g = generate_song(
             M["gen"], mel, course, level=int(level),
             density_bucket=COURSE_DENS[course], greedy=not sampling,
-            seed=0, device=DEVICE, plan=plan,
+            temperature=temperature, top_p=top_p, seed=0, device=DEVICE, plan=plan,
             on_progress=lambda d, t: P(0.32 + 0.58 * d / max(t, 1),
                                        f"Generating notes… {d}/{t} windows"))
         g = snap_chart(g, bpm)
@@ -360,7 +364,8 @@ def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_pla
         grid_info = " · grid: unreliable, barline anchoring off"
     else:
         grid_info = ""
-    info = (f"BPM {bpm:.1f} · {len(g['hits'])} notes · {len(g['spans'])} spans" + grid_info
+    sampling_info = f" · sampling T {temperature:.2f}, top-p {top_p:.2f}" if sampling else ""
+    info = (f"BPM {bpm:.1f} · {len(g['hits'])} notes · {len(g['spans'])} spans" + grid_info + sampling_info
             + (f" · plan: {sum(1 for b in plan if b[3]==1)} gaps, {sum(1 for b in plan if b[3]==2)} climax" if plan else ""))
     return tja_path, info, tja_img, audio_plan_img
 
@@ -385,13 +390,22 @@ with gr.Blocks(title="SoftChart — AI Taiko chart generator") as demo:
             with gr.Row():
                 use_beat = gr.Checkbox(label="Beat-anchor barlines", value=True)
                 sampling = gr.Checkbox(label="Sampling (diverse) vs greedy (best)", value=False)
+            with gr.Row():
+                temperature = gr.Slider(0.2, 1.2, value=DEFAULT_TEMPERATURE, step=0.05,
+                                        label="Temperature", interactive=False)
+                top_p = gr.Slider(0.5, 1.0, value=DEFAULT_TOP_P, step=0.01,
+                                  label="Top-p", interactive=False)
             btn = gr.Button("Generate chart", variant="primary")
         with gr.Column():
             out_tja = gr.File(label="Download .tja")
             out_info = gr.Textbox(label="Result", interactive=False)
             out_chart = gr.Image(label="TJA chart")
             out_img = gr.Image(label="Full-song mel + plan")
-    btn.click(generate, [audio, course, level, bpm, auto_plan_on, use_beat, use_planner, sampling],
+    sampling.change(
+        lambda enabled: (gr.update(interactive=enabled), gr.update(interactive=enabled)),
+        sampling, [temperature, top_p])
+    btn.click(generate, [audio, course, level, bpm, auto_plan_on, use_beat, use_planner,
+                         sampling, temperature, top_p],
               [out_tja, out_info, out_chart, out_img])
 
 if __name__ == "__main__":