Update app.py

app.py

@@ -56,23 +56,29 @@ Rules:
 
 Output only the final instruction in plain text and nothing else."""
 
-# Model repository IDs
-REPO_ID_4B = "diffusers-internal-dev/dummy-1015-4b"  # 4b model
-REPO_ID_9B = "diffusers-internal-dev/dummy-1015-9b"  # 9b model
+# Model repository IDs for 4B
+REPO_ID_REGULAR = "diffusers-internal-dev/dummy-1015-4b"
+REPO_ID_DISTILLED = "diffusers-internal-dev/dummy-1015-4b-distilled"
 
-# Load both models
-print("Loading 4B model...")
-pipe_4b = Flux2KleinPipeline.from_pretrained(REPO_ID_4B, torch_dtype=dtype)
-pipe_4b.to("cuda")
+# Load both 4B models
+print("Loading 4B Regular model...")
+pipe_regular = Flux2KleinPipeline.from_pretrained(REPO_ID_REGULAR, torch_dtype=dtype)
+pipe_regular.to("cuda")
 
-print("Loading 9B model...")
-pipe_9b = Flux2KleinPipeline.from_pretrained(REPO_ID_9B, torch_dtype=dtype)
-pipe_9b.to("cuda")
+print("Loading 4B Distilled model...")
+pipe_distilled = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
+pipe_distilled.to("cuda")
 
 # Dictionary for easy access
 pipes = {
-    "4B": pipe_4b,
-    "9B": pipe_9b,
+    "Distilled (4 steps)": pipe_distilled,
+    "Regular (30 steps)": pipe_regular,
+}
+
+# Default steps for each mode
+DEFAULT_STEPS = {
+    "Distilled (4 steps)": 4,
+    "Regular (30 steps)": 30,
 }
 
 
@@ -153,14 +159,19 @@ def update_dimensions_from_image(image_list):
     return new_width, new_height
 
 
+def update_steps_from_mode(mode_choice):
+    """Update the number of inference steps based on the selected mode."""
+    return DEFAULT_STEPS[mode_choice]
+
+
 @spaces.GPU(duration=85)
-def infer(prompt, input_images=None, model_choice="4B", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=4.0, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, input_images=None, mode_choice="Distilled (4 steps)", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=4.0, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    # Select the appropriate pipeline based on model choice
-    pipe = pipes[model_choice]
+    # Select the appropriate pipeline based on mode choice
+    pipe = pipes[mode_choice]
     
     # Prepare image list (convert None or empty gallery to None)
     image_list = None
@@ -178,7 +189,7 @@ def infer(prompt, input_images=None, model_choice="4B", seed=42, randomize_seed=
         print(f"Upsampled Prompt: {final_prompt}")
 
     # 2. Image Generation
-    progress(0.2, desc=f"Generating image with {model_choice} model...")
+    progress(0.2, desc=f"Generating image with 4B {mode_choice}...")
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
@@ -221,11 +232,11 @@ css = """
 }
 """
 
-with gr.Blocks() as demo:
+with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.2 [Klein]
-FLUX.2 [Klein] is a distilled model capable of generating, editing and combining images based on text instructions [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
+        gr.Markdown(f"""# FLUX.2 [Klein] - 4B (Apache 2.0)
+FLUX.2 [Klein] is a... [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
         """)
         with gr.Row():
             with gr.Column():
@@ -249,10 +260,10 @@ FLUX.2 [Klein] is a distilled model capable of generating, editing and combining
                         rows=1,
                     )
 
-                model_choice = gr.Radio(
-                    label="Model Size",
-                    choices=["4B", "9B"],
-                    value="4B",
+                mode_choice = gr.Radio(
+                    label="Mode",
+                    choices=["Distilled (4 steps)", "Regular (30 steps)"],
+                    value="Distilled (4 steps)",
                 )
                 
                 with gr.Accordion("Advanced Settings", open=False):
@@ -298,7 +309,7 @@ FLUX.2 [Klein] is a distilled model capable of generating, editing and combining
                             minimum=1,
                             maximum=100,
                             step=1,
-                            value=50,
+                            value=4,
                         )
                         
                         guidance_scale = gr.Slider(
@@ -338,12 +349,19 @@ FLUX.2 [Klein] is a distilled model capable of generating, editing and combining
         inputs=[input_images],
         outputs=[width, height]
     )
+    
+    # Auto-update steps when mode changes
+    mode_choice.change(
+        fn=update_steps_from_mode,
+        inputs=[mode_choice],
+        outputs=[num_inference_steps]
+    )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, input_images, model_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
+        inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
         outputs=[result, seed]
     )
 
-demo.launch(css=css)
+demo.launch()