UI: Citrus theme, image-first gallery, randomize seed, β=1.2 default, cached examples

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/chinese_vases/vase_03.jpg filter=lfs diff=lfs merge=lfs -text
+examples/pink_elephant/bank_0000.png filter=lfs diff=lfs merge=lfs -text
+examples/pink_elephant/bank_0001.png filter=lfs diff=lfs merge=lfs -text
+examples/pink_elephant/bank_0002.png filter=lfs diff=lfs merge=lfs -text
+examples/pink_elephant/bank_0003.png filter=lfs diff=lfs merge=lfs -text
+examples/pink_elephant/bank_0004.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -5,8 +5,8 @@ Code:  https://github.com/pedrocurvo/follow-the-mean
 
 The core idea: in flow matching, the velocity field is governed by an endpoint
 mean. Shift that endpoint mean toward a reference set and the flow follows it.
-This Space exposes the training-free RMG variant: pick a reference (a prompt
-that will be sampled M times, or a folder of images), then guide a frozen
+This Space exposes the training-free RMG variant: pick a reference (a folder
+of images or a prompt that will be sampled M times), then guide a frozen
 FLUX.2-klein generator with the empirical reference endpoint mean.
 """
 
@@ -30,6 +30,8 @@ import retrieval_guidance_core as poc
 MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"
 DEVICE = "cuda"
 DTYPE = torch.bfloat16
+MAX_SEED = 2_147_483_647
+HERE = Path(__file__).resolve().parent
 
 # Load the pipeline at module scope. ZeroGPU intercepts CUDA at import time,
 # but the weights still need to be resident when the @spaces.GPU function runs.
@@ -135,12 +137,29 @@ def _make_grid(images: List[Image.Image], cell: int = 256) -> Image.Image:
     return grid
 
 
+def _coerce_gallery_value(value) -> List[Image.Image]:
+    """Normalize a gr.Gallery value to a list of PIL images."""
+    if value is None:
+        return []
+    items = []
+    for entry in value:
+        if isinstance(entry, (list, tuple)):
+            entry = entry[0]
+        if isinstance(entry, dict):
+            entry = entry.get("image") or entry.get("name") or entry.get("path")
+        if isinstance(entry, Image.Image):
+            items.append(entry)
+        elif isinstance(entry, (str, Path)):
+            items.append(Image.open(entry))
+    return items
+
+
 @spaces.GPU(duration=300)
 def run_rmg(
     prompt: str,
     reference_mode: str,
     reference_prompt: str,
-    reference_images: Optional[List[Image.Image]],
+    reference_images,
     reference_size: int,
     num_inference_steps: int,
     guidance_strength: float,
@@ -150,17 +169,22 @@ def run_rmg(
     topk: int,
     height: int,
     width: int,
+    randomize_seed: bool,
     seed: int,
     progress=gr.Progress(track_tqdm=True),
 ):
     if not prompt or not prompt.strip():
         raise gr.Error("Please enter a prompt.")
 
-    reference_mode = reference_mode.lower()
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+
+    reference_mode = (reference_mode or "").lower()
     if reference_mode.startswith("image"):
-        if not reference_images:
-            raise gr.Error("Upload at least one reference image, or switch to prompt mode.")
-        ref_pils = [Image.open(p) if isinstance(p, str) else p for p in reference_images]
+        ref_pils = _coerce_gallery_value(reference_images)
+        if not ref_pils:
+            raise gr.Error("Add at least one reference image, or switch to prompt mode.")
         ref_pils = _prepare_image_reference(ref_pils, width, height)
     else:
         rp = (reference_prompt or "").strip()
@@ -220,7 +244,7 @@ def run_rmg(
     ).images[0]
 
     reference_grid = _make_grid(ref_pils)
-    return baseline, guided, reference_grid
+    return baseline, guided, reference_grid, seed
 
 
 DESCRIPTION = """
@@ -233,49 +257,48 @@ goes. Change the reference set, and the flow changes.
 
 * Paper: [Follow the Mean: Reference-Guided Flow Matching](https://arxiv.org/abs/2605.10302)
 * Code: [github.com/pedrocurvo/follow-the-mean](https://github.com/pedrocurvo/follow-the-mean)
+"""
 
-**Reference mode**
-* *Prompt* — sample N reference images from the given reference prompt.
-* *Images* — upload your own reference photos.
 
-Each run produces the **baseline** (prompt only), the **RMG-guided** image
-(same prompt, same base noise, plus reference-mean correction), and a grid of
-the references that drove the guidance.
-"""
+PINK_ELEPHANT_IMAGES = sorted(str(p) for p in (HERE / "examples" / "pink_elephant").glob("*.png"))
+CHINESE_VASE_IMAGES = sorted(
+    str(p) for p in (HERE / "examples" / "chinese_vases").glob("*.*")
+)
+
 
 EXAMPLES = [
     [
-        "an elephant in a jungle",
-        "Prompt",
-        "a pink elephant",
-        None,
-        4, 28, 0.5, "quadratic-decay", 0.15, 0.95, 4, 1024, 1024, 123,
-    ],
-    [
-        "a cat",
-        "Prompt",
-        "a van gogh painting",
-        None,
-        4, 28, 0.5, "quadratic-decay", 0.15, 0.95, 4, 1024, 1024, 123,
+        "an elephant in a jungle",         # prompt
+        "Images",                          # reference_mode
+        "",                                # reference_prompt (unused in image mode)
+        PINK_ELEPHANT_IMAGES,              # reference_images (gallery)
+        4,                                 # reference_size (unused in image mode)
+        28,                                # num_inference_steps
+        1.2,                               # guidance_strength
+        "quadratic-decay",                 # beta_schedule
+        0.15, 0.95,                        # start / end frac
+        4,                                 # topk
+        1024, 1024,                        # height, width
+        False, 123,                        # randomize_seed, seed
     ],
     [
-        "an animal in a savanna",
-        "Prompt",
-        "a giraffe",
-        None,
-        4, 28, 0.6, "quadratic-decay", 0.15, 0.95, 4, 1024, 1024, 7,
-    ],
-    [
-        "a house in a forest",
-        "Prompt",
-        "a pencil sketch",
-        None,
-        4, 28, 0.5, "bell", 0.10, 0.90, 4, 1024, 1024, 42,
+        "a capybara",
+        "Images",
+        "",
+        CHINESE_VASE_IMAGES,
+        4,
+        28,
+        1.25,
+        "quadratic-decay",
+        0.15, 0.95,
+        4,
+        1024, 1024,
+        False, 42,
     ],
 ]
 
 
-with gr.Blocks(title="Follow the Mean — FLUX.2", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Follow the Mean — FLUX.2", theme=gr.themes.Citrus()) as demo:
     gr.Markdown(DESCRIPTION)
 
     with gr.Row():
@@ -290,31 +313,35 @@ with gr.Blocks(title="Follow the Mean — FLUX.2", theme=gr.themes.Soft()) as de
             with gr.Group():
                 reference_mode = gr.Radio(
                     label="Reference mode",
-                    choices=["Prompt", "Images"],
-                    value="Prompt",
+                    choices=["Images", "Prompt"],
+                    value="Images",
+                )
+                reference_images = gr.Gallery(
+                    label="Reference images",
+                    columns=4,
+                    height=240,
+                    object_fit="cover",
+                    type="filepath",
+                    interactive=True,
+                    visible=True,
                 )
                 reference_prompt = gr.Textbox(
                     label="Reference prompt (used in Prompt mode)",
                     value="a pink elephant",
                     placeholder="The attribute / object / style to inject.",
                     lines=1,
-                )
-                reference_images = gr.File(
-                    label="Reference images (used in Images mode)",
-                    file_count="multiple",
-                    file_types=["image"],
-                    type="filepath",
                     visible=False,
                 )
                 reference_size = gr.Slider(
                     label="Reference set size (prompt mode)",
                     minimum=1, maximum=12, step=1, value=4,
+                    visible=False,
                 )
 
-            with gr.Accordion("Guidance", open=True):
+            with gr.Accordion("Guidance", open=False):
                 guidance_strength = gr.Slider(
                     label="Guidance strength (β scale)",
-                    minimum=0.0, maximum=2.0, step=0.05, value=0.5,
+                    minimum=0.0, maximum=2.0, step=0.05, value=1.2,
                 )
                 beta_schedule = gr.Radio(
                     label="β schedule",
@@ -340,7 +367,8 @@ with gr.Blocks(title="Follow the Mean — FLUX.2", theme=gr.themes.Soft()) as de
                 with gr.Row():
                     height = gr.Slider(label="Height", minimum=512, maximum=1280, step=64, value=1024)
                     width = gr.Slider(label="Width", minimum=512, maximum=1280, step=64, value=1024)
-                seed = gr.Slider(label="Seed", minimum=0, maximum=2_147_483_647, step=1, value=123)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=123)
 
             run_btn = gr.Button("Generate", variant="primary")
 
@@ -351,7 +379,7 @@ with gr.Blocks(title="Follow the Mean — FLUX.2", theme=gr.themes.Soft()) as de
             reference_out = gr.Image(label="Reference set", interactive=False)
 
     def _toggle_reference_inputs(mode):
-        is_images = mode.lower().startswith("image")
+        is_images = (mode or "").lower().startswith("image")
         return (
             gr.update(visible=is_images),       # reference_images
             gr.update(visible=not is_images),   # reference_prompt
@@ -364,48 +392,34 @@ with gr.Blocks(title="Follow the Mean — FLUX.2", theme=gr.themes.Soft()) as de
         outputs=[reference_images, reference_prompt, reference_size],
     )
 
-    run_btn.click(
-        run_rmg,
-        inputs=[
-            prompt,
-            reference_mode,
-            reference_prompt,
-            reference_images,
-            reference_size,
-            num_inference_steps,
-            guidance_strength,
-            beta_schedule,
-            guidance_start_frac,
-            guidance_end_frac,
-            topk,
-            height,
-            width,
-            seed,
-        ],
-        outputs=[baseline_out, guided_out, reference_out],
-    )
+    inputs_list = [
+        prompt,
+        reference_mode,
+        reference_prompt,
+        reference_images,
+        reference_size,
+        num_inference_steps,
+        guidance_strength,
+        beta_schedule,
+        guidance_start_frac,
+        guidance_end_frac,
+        topk,
+        height,
+        width,
+        randomize_seed,
+        seed,
+    ]
+    outputs_list = [baseline_out, guided_out, reference_out, seed]
+
+    run_btn.click(run_rmg, inputs=inputs_list, outputs=outputs_list)
 
     gr.Examples(
         examples=EXAMPLES,
-        inputs=[
-            prompt,
-            reference_mode,
-            reference_prompt,
-            reference_images,
-            reference_size,
-            num_inference_steps,
-            guidance_strength,
-            beta_schedule,
-            guidance_start_frac,
-            guidance_end_frac,
-            topk,
-            height,
-            width,
-            seed,
-        ],
-        outputs=[baseline_out, guided_out, reference_out],
+        inputs=inputs_list,
+        outputs=outputs_list,
         fn=run_rmg,
-        cache_examples=False,
+        cache_examples=True,
+        cache_mode="lazy",
     )