change guessing to use buttons

app.py

@@ -1,9 +1,6 @@
 import torch
 import gradio as gr
 from diffusers import FluxPipeline, FluxTransformer2DModel
-from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
-from transformers import T5EncoderModel
-from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 import gc
 import random
 from PIL import Image
@@ -131,6 +128,9 @@ def generate_images(prompt, quantization_choice, progress=gr.Progress(track_tqdm
 
     current_pipe = None # Keep track of the current pipe for cleanup
 
+    seed = random.getrandbits(64)
+    print(f"Using seed: {seed}")
+
     for i, (label, load_func) in enumerate(model_configs):
         progress(i / len(model_configs), desc=f"Loading {label} model...")
         print(f"\n--- Loading {label} Model ---")
@@ -143,8 +143,9 @@ def generate_images(prompt, quantization_choice, progress=gr.Progress(track_tqdm
             progress((i + 0.5) / len(model_configs), desc=f"Generating with {label} model...")
             print(f"--- Generating with {label} Model ---")
             gen_start_time = time.time()
-            image_list = current_pipe(**pipe_kwargs, generator=torch.manual_seed(GENERATION_SEED)).images
+            image_list = current_pipe(**pipe_kwargs, generator=torch.manual_seed(seed)).images
             image = image_list[0]
+            # image.save(f"{load_start_time}.png")
             gen_end_time = time.time()
             results.append({"label": label, "image": image})
             print(f"--- Finished Generation with {label} Model in {gen_end_time - gen_start_time:.2f} seconds ---")
@@ -219,11 +220,11 @@ with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as d
     gr.Markdown(
         "Compare the original FLUX.1-dev (BF16) model against a quantized version (4-bit or 8-bit). "
         "Enter a prompt, choose the quantization method, and generate two images. "
-        "The images will be shuffled. Can you guess which one used quantization?"
+        "The images will be shuffled, can you spot which one was quantized?"
     )
 
     with gr.Row():
-        prompt_input = gr.Textbox(label="Enter Prompt", placeholder="e.g., A photorealistic portrait of an astronaut on Mars", scale=3)
+        prompt_input = gr.Textbox(label="Enter Prompt", scale=3)
         quantization_choice_radio = gr.Radio(
             choices=["8-bit", "4-bit"],
             label="Select Quantization",
@@ -243,20 +244,8 @@ with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as d
 
     gr.Markdown("### Which image used the selected quantization method?")
     with gr.Row():
-         # Centered guess radio and submit button
-         with gr.Column(scale=1): # Dummy column for spacing
-             pass
-         with gr.Column(scale=2): # Column for the radio button
-            guess_radio = gr.Radio(
-                choices=[],
-                label="Your Guess",
-                info="Select the image you believe was generated with the quantized model.",
-                interactive=False # Disabled until images are generated
-            )
-         with gr.Column(scale=1): # Column for the button
-            submit_guess_button = gr.Button("Submit Guess")
-         with gr.Column(scale=1): # Dummy column for spacing
-             pass
+        image1_btn = gr.Button("Image 1")
+        image2_btn = gr.Button("Image 2")
 
     feedback_box = gr.Textbox(label="Feedback", interactive=False, lines=1)
 
@@ -267,20 +256,18 @@ with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as d
     generate_button.click(
         fn=generate_images,
         inputs=[prompt_input, quantization_choice_radio],
-        outputs=[output_gallery, correct_mapping_state, feedback_box, guess_radio]
-    ).then(
-        lambda: "", # Clear feedback box on new generation
-        outputs=[feedback_box]
-    )
+        outputs=[output_gallery, correct_mapping_state],
+    ).then(lambda: "", outputs=[feedback_box])  # clear feedback on new run
 
+    # 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)
 
-    submit_guess_button.click(
-        fn=check_guess,
-        inputs=[guess_radio, correct_mapping_state], # Pass the selected guess and the state
-        outputs=[feedback_box]
-    )
+    image1_btn.click(choose_img1, inputs=[correct_mapping_state], outputs=[feedback_box])
+    image2_btn.click(choose_img2, inputs=[correct_mapping_state], outputs=[feedback_box])
 
 if __name__ == "__main__":
-    # queue()
-    # demo.queue().launch()
-    demo.launch(share=True)
+    demo.launch(share=True)
+    demo.launch()