add examples

.gitattributes

@@ -33,3 +33,7 @@ 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/astronauts_bnb_8bit.png filter=lfs diff=lfs merge=lfs -text
+examples/astronauts_seed_6456306350371904162.png filter=lfs diff=lfs merge=lfs -text
+examples/watercolor_cat_bnb_8bit.png filter=lfs diff=lfs merge=lfs -text
+examples/watercolor_cat_seed_14269059182221286790.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -3,6 +3,8 @@ import gradio as gr
 from diffusers import FluxPipeline, FluxTransformer2DModel
 import gc
 import random
+import glob
+from pathlib import Path
 from PIL import Image
 import os
 import time
@@ -214,6 +216,33 @@ def check_guess(user_guess, correct_mapping_state):
 
     return feedback
 
+EXAMPLE_DIR = Path(__file__).parent / "examples"
+EXAMPLES = [
+    {
+        "prompt": "A photorealistic portrait of an astronaut on Mars",
+        "files": ["astronauts_seed_6456306350371904162.png", "astronauts_bnb_8bit.png"],
+        "quantized_idx": 1,          # which of the two files is the quantized result
+    },
+    {
+        "prompt": "Water-color painting of a cat wearing sunglasses",
+        "files": ["watercolor_cat_bnb_8bit.png", "watercolor_cat_seed_14269059182221286790.png"],
+        "quantized_idx": 0,
+    },
+    # {
+    #     "prompt": "Neo-tokyo cyberpunk cityscape at night, rain-soaked streets, 8-K",
+    #     "files": ["cyber_city_q.jpg", "cyber_city.jpg"],
+    #     "quantized_idx": 0,
+    # },
+]
+
+def load_example(idx):
+    """Return [(PIL.Image, caption)...], mapping dict, and feedback string"""
+    ex = EXAMPLES[idx]
+    imgs = [Image.open(EXAMPLE_DIR / f) for f in ex["files"]]
+    gallery_items = [(img, f"Image {i+1}") for i, img in enumerate(imgs)]
+    mapping = {i: ("Quantized" if i == ex["quantized_idx"] else "Original")
+               for i in range(2)}
+    return gallery_items, mapping, f"{ex['prompt']}"
 
 with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# FLUX Model Quantization Challenge")
@@ -223,6 +252,13 @@ with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as d
         "The images will be shuffled, can you spot which one was quantized?"
     )
 
+    gr.Markdown("### Examples")
+    ex_selector = gr.Radio(
+        choices=[f"Example {i+1}" for i in range(len(EXAMPLES))],
+        label="Choose an example prompt",
+        interactive=True,
+    )
+    gr.Markdown("### …or create your own comparison")
     with gr.Row():
         prompt_input = gr.Textbox(label="Enter Prompt", scale=3)
         quantization_choice_radio = gr.Radio(
@@ -239,7 +275,7 @@ with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as d
         height=512,
         object_fit="contain",
         allow_preview=True,
-        show_label=True, # Shows "Image 1", "Image 2" captions we provide
+        show_label=True,
     )
 
     gr.Markdown("### Which image used the selected quantization method?")
@@ -253,21 +289,44 @@ with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as d
     # e.g., {0: 'Original', 1: 'Quantized (8-bit)'} or {0: 'Quantized (4-bit)', 1: 'Original'}
     correct_mapping_state = gr.State({})
 
+    def _load_example(sel):
+        idx = int(sel.split()[-1]) - 1
+        return load_example(idx)
+
+    ex_selector.change(
+        fn=_load_example,
+        inputs=ex_selector,
+        outputs=[output_gallery, correct_mapping_state, prompt_input],
+    ).then(
+        lambda: (gr.update(interactive=True), gr.update(interactive=True)),
+        outputs=[image1_btn, image2_btn],
+    )
+
     generate_button.click(
         fn=generate_images,
         inputs=[prompt_input, quantization_choice_radio],
-        outputs=[output_gallery, correct_mapping_state],
-    ).then(lambda: "", outputs=[feedback_box])  # clear feedback on new run
+        outputs=[output_gallery, correct_mapping_state] #, feedback_box],
+    ).then(
+        lambda: (gr.update(interactive=True),
+                 gr.update(interactive=True),
+                 ""),               # clear feedback
+        outputs=[image1_btn, image2_btn, feedback_box],
+    )
 
-    # helper wrappers so we can supply the fixed choice string
-    def choose_img1(mapping): 
-        return check_guess("Image 1", mapping)
-    def choose_img2(mapping): 
-        return check_guess("Image 2", mapping)
+    def choose(choice_string, mapping):
+        feedback = check_guess(choice_string, mapping)
+        return feedback, gr.update(interactive=False), gr.update(interactive=False)
 
-    image1_btn.click(choose_img1, inputs=[correct_mapping_state], outputs=[feedback_box])
-    image2_btn.click(choose_img2, inputs=[correct_mapping_state], outputs=[feedback_box])
+    image1_btn.click(
+        fn=lambda mapping: choose("Image 1", mapping),
+        inputs=[correct_mapping_state],
+        outputs=[feedback_box, image1_btn, image2_btn],
+    )
+    image2_btn.click(
+        fn=lambda mapping: choose("Image 2", mapping),
+        inputs=[correct_mapping_state],
+        outputs=[feedback_box, image1_btn, image2_btn],
+    )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
-    demo.launch()
+    demo.launch(share=True)