Migrate from FLUX.1-schnell to FLUX.2-klein-base-4B

Dockerfile

@@ -5,7 +5,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     PORT=7860 \
     HOST=0.0.0.0 \
     GRADIO_SERVER_NAME=0.0.0.0 \
-    P5_FLUX_MODEL_DIR=/app/models/FLUX.1-schnell
+    P5_FLUX_MODEL_DIR=/app/models/FLUX.2-klein-base-4B
 
 RUN apt-get update \
     && apt-get install -y --no-install-recommends git ca-certificates libcurl4 libgomp1 libstdc++6 \
@@ -28,7 +28,7 @@ ENV HF_TOKEN=${HF_TOKEN}
 RUN pip install huggingface_hub hf_transfer
 RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=false \
     if [ -f /run/secrets/HF_TOKEN ]; then export HF_TOKEN=$(cat /run/secrets/HF_TOKEN); fi && \
-    python scripts/download_model.py black-forest-labs/FLUX.1-schnell models/FLUX.1-schnell && \
+    python scripts/download_model.py black-forest-labs/FLUX.2-klein-base-4B models/FLUX.2-klein-base-4B && \
     python scripts/download_model.py openbmb/MiniCPM5-1B models/MiniCPM5-1B
 
 EXPOSE 7860

configs/model_registry.yaml

@@ -17,7 +17,7 @@ projects:
   p5:
     image_gen:
       component: image_gen
-      model_id: black-forest-labs/FLUX.1-schnell
+      model_id: black-forest-labs/FLUX.2-klein-base-4B
       license: Apache-2.0
       usage_notes: Primary tile renderer for quilt art; runtime now fails clearly if no local checkpoint is mounted.
       runtime: heuristic+hf

configs/sponsor_model_policy.yaml

@@ -12,4 +12,4 @@ required_models:
     p5:
       image_gen:
         component: image_gen
-        model_id: black-forest-labs/FLUX.1-schnell
+        model_id: black-forest-labs/FLUX.2-klein-base-4B

requirements.txt

@@ -4,7 +4,7 @@ gradio>=5.0,<6
 Pillow>=10.4,<11
 PyYAML>=6.0,<7
 torch
-diffusers>=0.30.0
+diffusers>=0.31.0
 transformers>=4.43.0
 accelerate>=0.31.0
 sentencepiece

run_local.sh

@@ -5,7 +5,7 @@ cd "$ROOT"
 
 export P5_MEMORY_QUILT_FALLBACK_MODEL_PATH=$(ls -d /root/docker-data/hf_cache/hub/models--openbmb--MiniCPM5-1B/snapshots/* | head -n 1)
 export P5_MEMORY_QUILT_PRIMARY_MODEL_PATH=$(ls -d /root/docker-data/hf_cache/hub/models--JetBrains--Mellum2-12B-A2.5B-Instruct/snapshots/* 2>/dev/null | head -n 1 || echo "$P5_MEMORY_QUILT_FALLBACK_MODEL_PATH")
-export P5_FLUX_MODEL_DIR=$(ls -d /root/docker-data/hf_cache/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/* | head -n 1)
+export P5_FLUX_MODEL_DIR=$(ls -d /root/docker-data/hf_cache/hub/models--black-forest-labs--FLUX.2-klein-base-4B/snapshots/* | head -n 1)
 
 if [ ! -d "/tmp/p5_venv" ]; then
     python3 -m venv /tmp/p5_venv

src/apps/p5_memory_quilt/flux_backend.py

@@ -16,7 +16,7 @@ import time
 
 @dataclass(frozen=True)
 class FluxConfig:
-    model_id: str = "black-forest-labs/FLUX.1-schnell"
+    model_id: str = "black-forest-labs/FLUX.2-klein-base-4B"
     model_dir: str = ""
     steps: int = 4
     guidance_scale: float | None = None
@@ -26,7 +26,7 @@ class FluxConfig:
 
 
 def _load_config(size: int = 512) -> FluxConfig:
-    model_id = os.environ.get("P5_FLUX_MODEL_ID", "black-forest-labs/FLUX.1-schnell").strip()
+    model_id = os.environ.get("P5_FLUX_MODEL_ID", "black-forest-labs/FLUX.2-klein-base-4B").strip()
     model_dir = os.environ.get("P5_FLUX_MODEL_DIR", "").strip()
     steps = int(os.environ.get("P5_FLUX_STEPS", "4"))
     guidance_raw = os.environ.get("P5_FLUX_GUIDANCE", "").strip()
@@ -78,8 +78,10 @@ def try_generate(prompt: str, size: int = 512):
 
     model_ref: str = cfg.model_dir
 
-    # Prefer the explicit FluxPipeline if present; otherwise fall back to DiffusionPipeline.
-    PipelineCls = getattr(diffusers, "FluxPipeline", None)
+    # Prefer the explicit Flux2KleinPipeline if present; otherwise fall back to FluxPipeline or DiffusionPipeline.
+    PipelineCls = getattr(diffusers, "Flux2KleinPipeline", None)
+    if PipelineCls is None:
+        PipelineCls = getattr(diffusers, "FluxPipeline", None)
     if PipelineCls is None:
         from diffusers import DiffusionPipeline as PipelineCls  # type: ignore
 

tests/test_sponsor_model_policy.py

@@ -33,7 +33,7 @@ def test_sponsor_model_checker_fails_without_waiver_for_mismatch(tmp_path: Path)
     completed = run_checker("--model-registry", str(mismatch_path), "--policy", str(POLICY))
     assert completed.returncode == 1
     assert "sponsor model mismatch" in completed.stderr.lower()
-    assert "black-forest-labs/FLUX.1-schnell" in completed.stderr
+    assert "black-forest-labs/FLUX.2-klein-base-4B" in completed.stderr
 
 
 def test_sponsor_model_checker_passes_with_valid_waiver_and_warns(tmp_path: Path) -> None: