Update app.py

app.py

@@ -56,16 +56,16 @@ Rules:
 
 Output only the final instruction in plain text and nothing else."""
 
-# Model repository IDs for 4B
-REPO_ID_REGULAR = "diffusers-internal-dev/dummy-1015-4b"
-REPO_ID_DISTILLED = "diffusers-internal-dev/dummy-1015-4b-distilled"
+# Model repository IDs for 9B
+REPO_ID_REGULAR = "diffusers-internal-dev/dummy-1015-9b"
+REPO_ID_DISTILLED = "diffusers-internal-dev/dummy-1015-9b-distilled"
 
-# Load both 4B models
-print("Loading 4B Regular model...")
+# Load both 9B models
+print("Loading 9B Regular model...")
 pipe_regular = Flux2KleinPipeline.from_pretrained(REPO_ID_REGULAR, torch_dtype=dtype)
 pipe_regular.to("cuda")
 
-print("Loading 4B Distilled model...")
+print("Loading 9B Distilled model...")
 pipe_distilled = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
 pipe_distilled.to("cuda")
 
@@ -189,7 +189,7 @@ def infer(prompt, input_images=None, mode_choice="Distilled (4 steps)", seed=42,
         print(f"Upsampled Prompt: {final_prompt}")
 
     # 2. Image Generation
-    progress(0.2, desc=f"Generating image with 4B {mode_choice}...")
+    progress(0.2, desc=f"Generating image with 9B {mode_choice}...")
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
@@ -235,8 +235,8 @@ css = """
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        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)]
+        gr.Markdown(f"""# FLUX.2 [Klein] - 9B
+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)]
         """)
         with gr.Row():
             with gr.Column():