Restructure: per-prompt blocks with Baseline then Ours (1 image + 3 audios each)

app.py

@@ -14,22 +14,19 @@ TOP1_ID = SAMPLE_IDS[0] if SAMPLE_IDS else None
 
 
 def build_noise_demo(sample_id: str | None):
-    """Returns (baseline, nn1, nn2, nn3) each: (spec, bg_wav, fg_wav, m_wav); nn1/nn2/nn3 also have prompt."""
+    """Returns for each of 3 blocks: prompt_md, baseline (spec, bg, fg, m), method (spec, bg, fg, m)."""
     if not sample_id:
-        return (None,) * 19
+        return (None,) * 27
     data = get_noise_demo_paths(sample_id)
     out = []
-    for key in ("baseline", "nn1", "nn2", "nn3"):
-        block = data.get(key, {})
-        out.extend([
-            block.get("spec"),
-            block.get("bg_wav"),
-            block.get("fg_wav"),
-            block.get("m_wav"),
-        ])
-        if key.startswith("nn"):
-            prompt = block.get("prompt", "") or ""
-            out.append(f"**Prompt:** {prompt}" if prompt else "")
+    for i in range(1, 4):
+        block = data.get(f"block{i}", {})
+        prompt = block.get("prompt", "") or ""
+        out.append(f"**Prompt:** {prompt}" if prompt else "")
+        bl = block.get("baseline", {})
+        out.extend([bl.get("spec"), bl.get("bg_wav"), bl.get("fg_wav"), bl.get("m_wav")])
+        nn = block.get("nn", {})
+        out.extend([nn.get("spec"), nn.get("bg_wav"), nn.get("fg_wav"), nn.get("m_wav")])
     return tuple(out)
 
 
@@ -54,37 +51,38 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
         label="Noise (ID)",
     )
 
-    def block_ui(title: str, with_prompt: bool = False):
+    def prompt_block_ui(block_label: str):
+        """One block: prompt text, then Baseline (img + 3 audio), then Ours (img + 3 audio)."""
         with gr.Group():
-            gr.Markdown(f"### {title}")
-            if with_prompt:
-                prompt_md = gr.Markdown(value="", elem_id=f"{title}_prompt")
-            img = gr.Image(label=f"{title}", show_label=True)
+            gr.Markdown(f"### {block_label}")
+            prompt_md = gr.Markdown(value="")
+            gr.Markdown("**Baseline**")
             with gr.Row():
-                abg = gr.Audio(label="BG", show_label=True)
-                afg = gr.Audio(label="FG", show_label=True)
-                am = gr.Audio(label="Mix", show_label=True)
-        if with_prompt:
-            return prompt_md, img, abg, afg, am
-        return img, abg, afg, am
-
-    # Baseline
-    bl_img, bl_bg, bl_fg, bl_m = block_ui("Baseline")
-    gr.Markdown("---")
-    # NN1
-    nn1_prompt, nn1_img, nn1_bg, nn1_fg, nn1_m = block_ui("NN1", with_prompt=True)
+                bl_img = gr.Image(label="Baseline", show_label=True)
+                bl_bg = gr.Audio(label="BG", show_label=True)
+                bl_fg = gr.Audio(label="FG", show_label=True)
+                bl_m = gr.Audio(label="Mix", show_label=True)
+            gr.Markdown("**Ours**")
+            with gr.Row():
+                nn_img = gr.Image(label="Ours", show_label=True)
+                nn_bg = gr.Audio(label="BG", show_label=True)
+                nn_fg = gr.Audio(label="FG", show_label=True)
+                nn_m = gr.Audio(label="Mix", show_label=True)
+        return prompt_md, bl_img, bl_bg, bl_fg, bl_m, nn_img, nn_bg, nn_fg, nn_m
+
+    # Block 1: Prompt 1 -> Baseline -> Ours
+    p1_md, bl1_img, bl1_bg, bl1_fg, bl1_m, nn1_img, nn1_bg, nn1_fg, nn1_m = prompt_block_ui("Prompt 1")
     gr.Markdown("---")
-    # NN2
-    nn2_prompt, nn2_img, nn2_bg, nn2_fg, nn2_m = block_ui("NN2", with_prompt=True)
+    # Block 2
+    p2_md, bl2_img, bl2_bg, bl2_fg, bl2_m, nn2_img, nn2_bg, nn2_fg, nn2_m = prompt_block_ui("Prompt 2")
     gr.Markdown("---")
-    # NN3
-    nn3_prompt, nn3_img, nn3_bg, nn3_fg, nn3_m = block_ui("NN3", with_prompt=True)
+    # Block 3
+    p3_md, bl3_img, bl3_bg, bl3_fg, bl3_m, nn3_img, nn3_bg, nn3_fg, nn3_m = prompt_block_ui("Prompt 3")
 
     all_outputs = [
-        bl_img, bl_bg, bl_fg, bl_m,
-        nn1_img, nn1_bg, nn1_fg, nn1_m, nn1_prompt,
-        nn2_img, nn2_bg, nn2_fg, nn2_m, nn2_prompt,
-        nn3_img, nn3_bg, nn3_fg, nn3_m, nn3_prompt,
+        p1_md, bl1_img, bl1_bg, bl1_fg, bl1_m, nn1_img, nn1_bg, nn1_fg, nn1_m,
+        p2_md, bl2_img, bl2_bg, bl2_fg, bl2_m, nn2_img, nn2_bg, nn2_fg, nn2_m,
+        p3_md, bl3_img, bl3_bg, bl3_fg, bl3_m, nn3_img, nn3_bg, nn3_fg, nn3_m,
     ]
 
     def on_noise_select(sid):

dataset_loader.py

@@ -103,10 +103,8 @@ def get_nn_demo_paths(bid: str) -> dict:
 
 def get_noise_demo_paths(bid: str) -> dict:
     """
-    One block per method: baseline, nn1, nn2, nn3.
-    Each block: one combined image (spec) + 3 audios (bg_wav, fg_wav, m_wav).
-    nn1/nn2/nn3 also include "prompt" (retrieved text). 
-    Returns { "baseline": {spec, bg_wav, fg_wav, m_wav}, "nn1": {..., prompt}, ... }.
+    One block per prompt (1, 2, 3): each has prompt text, baseline (spec + 3 wavs), and our method (spec + 3 wavs).
+    Returns { "block1": {prompt, baseline: {...}, nn: {...}}, "block2": ..., "block3": ... }.
     """
     inner = f"{ROOT_PREFIX}{bid}/{bid}"
     files = _find_files(inner)
@@ -141,22 +139,34 @@ def get_noise_demo_paths(bid: str) -> dict:
             "m_wav": _download_file(m) if m else None,
         }
 
-    # Baseline: use first generated_baseline_* (01 or 02 etc.)
-    baseline_prefix = None
+    # Find baseline folder names generated_baseline_01_*, 02_*, 03_*
+    seen = set()
+    baseline_folders = []
     for f in baseline_files:
-        if "/baseline/generated_baseline_" in f:
-            parts = f.replace(baseline_inner + "/", "").split("/")
-            if parts and parts[0].startswith("generated_baseline_"):
-                baseline_prefix = baseline_inner + "/" + parts[0]
-                break
-    baseline_block = collect_block(baseline_files, baseline_prefix) if baseline_prefix else {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
+        parts = f.replace(baseline_inner + "/", "").split("/")
+        if parts and parts[0].startswith("generated_baseline_") and parts[0] not in seen:
+            seen.add(parts[0])
+            baseline_folders.append((parts[0], baseline_inner + "/" + parts[0]))
+    baseline_folders.sort(key=lambda x: x[0])
 
-    # NN1, NN2, NN3: generated_01_, generated_02_, generated_03_ + prompt from retrieval JSON
-    result = {"baseline": baseline_block}
+    result = {}
     for i in range(1, 4):
+        prompt_text = prompts[i - 1].get("prompt", "") if i <= len(prompts) else ""
+        # Baseline for this prompt: i-th baseline folder (01, 02, 03)
+        bl_prefix = f"generated_baseline_{i:02d}_"
+        baseline_block = {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
+        for folder_name, full_prefix in baseline_folders:
+            if folder_name.startswith(bl_prefix):
+                baseline_block = collect_block(baseline_files, full_prefix)
+                break
+        # Our method: generated_0{i}_*
         rel_prefix = f"generated_{i:02d}_"
         nn_files = [f for f in files if f.replace(inner + "/", "").startswith(rel_prefix)]
-        block = collect_block(nn_files, rel_prefix)
-        block["prompt"] = prompts[i - 1].get("prompt", "") if i <= len(prompts) else ""
-        result[f"nn{i}"] = block
+        nn_block = collect_block(nn_files, rel_prefix)
+        nn_block["prompt"] = prompt_text
+        result[f"block{i}"] = {
+            "prompt": prompt_text,
+            "baseline": baseline_block,
+            "nn": nn_block,
+        }
     return result