update app

app.py

@@ -10,8 +10,9 @@ import random
 import base64
 from io import BytesIO
 from PIL import Image
+from typing import Tuple
 
-from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
+from diffusers import Flux2KleinPipeline
 
 MAX_SEED = np.iinfo(np.int32).max
 LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
@@ -21,19 +22,12 @@ dtype = torch.bfloat16
 
 print("Using device:", device)
 
-print("Loading Small Decoder VAE...")
-vae_small = AutoencoderKLFlux2.from_pretrained(
-    "black-forest-labs/FLUX.2-small-decoder",
-    torch_dtype=dtype,
-).to(device)
-
-print("Loading 4B Distilled model (Small Decoder VAE)...")
+print("Loading FLUX.2 Klein 9B model base...")
 pipe = Flux2KleinPipeline.from_pretrained(
-    "black-forest-labs/FLUX.2-klein-4B",
-    vae=vae_small,
+    "black-forest-labs/FLUX.2-klein-9B",
     torch_dtype=dtype,
 ).to(device)
-print("Pipeline loaded directly to CUDA.")
+print("Base Model loaded successfully.")
 
 def pil_to_b64_png(image: Image.Image) -> str:
     buf = BytesIO()
@@ -56,7 +50,7 @@ def infer(
     steps: int,
     guidance_scale: float,
 ) -> dict:
-    """Generates an image from a scribble using FLUX.2 Klein 4B."""
+    """Generates an image from a scribble using FLUX.2 Klein 9B."""
     gc.collect()
     torch.cuda.empty_cache()