Apply the real weights

app.py

@@ -1,15 +1,6 @@
 import os
 import subprocess
 import sys
-
-subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "diffusers"], check=False)
-subprocess.check_call(["unzip", "-o", "diffusers.zip", "-d", "diffusers_local"])
-subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", "./diffusers_local"])
-
-subprocess.check_call(["pwd"])
-subprocess.check_call(["ls", "-la"])
-
-
 import io
 import gradio as gr
 import numpy as np
@@ -57,8 +48,8 @@ Rules:
 Output only the final instruction in plain text and nothing else."""
 
 # Model repository IDs for 9B
-REPO_ID_REGULAR = "diffusers-internal-dev/dummy-1015-9b"
-REPO_ID_DISTILLED = "diffusers-internal-dev/dummy-1015-9b-distilled"
+REPO_ID_REGULAR = "black-forest-labs/FLUX.2-klein-base-9B"
+REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-9B"
 
 # Load both 9B models
 print("Loading 9B Regular model...")
@@ -72,18 +63,18 @@ pipe_distilled.to("cuda")
 # Dictionary for easy access
 pipes = {
     "Distilled (4 steps)": pipe_distilled,
-    "Regular (30 steps)": pipe_regular,
+    "Base (50 steps)": pipe_regular,
 }
 
 # Default steps for each mode
 DEFAULT_STEPS = {
     "Distilled (4 steps)": 4,
-    "Regular (30 steps)": 30,
+    "Base (50 steps)": 50,
 }
 
 DEFAULT_CFG = {
     "Distilled (4 steps)": 1.0,
-    "Regular (30 steps)": 4.0,
+    "Base (50 steps)": 4.0,
 }
 
 def image_to_data_uri(img):
@@ -266,7 +257,7 @@ FLUX.2 [Klein] is a distilled model capable of generating, editing and combining
 
                 mode_choice = gr.Radio(
                     label="Mode",
-                    choices=["Distilled (4 steps)", "Regular (30 steps)"],
+                    choices=["Distilled (4 steps)", "Base (50 steps)"],
                     value="Distilled (4 steps)",
                 )
                 

requirements.txt

@@ -1,3 +1,4 @@
+git+https://github.com/diffusers.git
 transformers
 accelerate
 safetensors