Update app.py

app.py

@@ -56,13 +56,24 @@ Rules:
 
 Output only the final instruction in plain text and nothing else."""
 
-# Load model - using Flux2KleinPipeline with built-in text encoder
-repo_id = "diffusers-internal-dev/dummy-1015-4b"  # 4b model
-# repo_id = "diffusers-internal-dev/dummy-1015-9b"  # 9b model (alternative)
-
-pipe = Flux2KleinPipeline.from_pretrained(repo_id, torch_dtype=dtype)
-pipe.to("cuda")
-#pipe.enable_model_cpu_offload()
+# 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
+
+# Load both models
+print("Loading 4B model...")
+pipe_4b = Flux2KleinPipeline.from_pretrained(REPO_ID_4B, torch_dtype=dtype)
+pipe_4b.to("cuda")
+
+print("Loading 9B model...")
+pipe_9b = Flux2KleinPipeline.from_pretrained(REPO_ID_9B, torch_dtype=dtype)
+pipe_9b.to("cuda")
+
+# Dictionary for easy access
+pipes = {
+    "4B": pipe_4b,
+    "9B": pipe_9b,
+}
 
 
 def image_to_data_uri(img):
@@ -143,11 +154,14 @@ def update_dimensions_from_image(image_list):
 
 
 @spaces.GPU(duration=85)
-def infer(prompt, input_images=None, 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, 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)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
+    # Select the appropriate pipeline based on model choice
+    pipe = pipes[model_choice]
+    
     # Prepare image list (convert None or empty gallery to None)
     image_list = None
     if input_images is not None and len(input_images) > 0:
@@ -164,7 +178,7 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         print(f"Upsampled Prompt: {final_prompt}")
 
     # 2. Image Generation
-    progress(0.2, desc="Generating image...")
+    progress(0.2, desc=f"Generating image with {model_choice} model...")
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
@@ -236,6 +250,13 @@ FLUX.2 Klein [dev] is a distilled model capable of generating, editing and combi
                     )
                 
                 with gr.Accordion("Advanced Settings", open=False):
+                    model_choice = gr.Radio(
+                        label="Model Size",
+                        choices=["4B", "9B"],
+                        value="4B",
+                        info="Choose between the 4B (faster) or 9B (higher quality) model"
+                    )
+                    
                     prompt_upsampling = gr.Checkbox(
                         label="Prompt Upsampling",
                         value=True,
@@ -321,7 +342,7 @@ FLUX.2 Klein [dev] is a distilled model capable of generating, editing and combi
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
+        inputs=[prompt, input_images, model_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
         outputs=[result, seed]
     )