feat: port img2img watercolour develop to ZeroGPU (staged LoRA)

app.py

@@ -1274,6 +1274,12 @@ def client_config():
         "asrProfiles": list(ASR_CHUNK_PROFILES.keys()),
         "asrDefaultProfile": ASR_DEFAULT_PROFILE,
         "captionModel": AUTOCAPTION_MODEL_ID,
+        # "Proceed to Development" watercolour: the endpoint always exists; this
+        # only reports whether torch/diffusers are present (a CUDA-free check —
+        # the LoRA is optional, stage 1 runs base img2img). The ~14 GB pipeline
+        # still builds lazily on first use inside the GPU worker.
+        "img2imgEnabled": _img2img_prep(),
+        "img2imgModel": IMG2IMG_MODEL_ID,
         "phoenix": _get_tracer() is not None,
     }
 
@@ -1411,6 +1417,240 @@ async def nsfw_rest(request: Request):
         return {"score": 0.0, "blocked": False, "threshold": NSFW_THRESHOLD, "available": available, "error": str(exc)}
 
 
+# --------------------------------------------------------------------------- #
+# Develop → watercolour via Flux.2-klein + a scene LoRA (in-process img2img).
+# This is what "Proceed to Development" triggers: each photo is sent here and
+# painted into a watercolour by Flux2KleinPipeline (the source is VAE-encoded
+# and used as reference conditioning while the LoRA steers the watercolour
+# style).
+#
+# ZeroGPU port notes (mirrors the caption/NSFW/ASR helpers in this file):
+#   * The pipeline is ~14 GB and lives ONLY on the GPU, so it must be built
+#     INSIDE a @spaces.GPU context — never in the main process, where CUDA must
+#     stay un-initialized. `_img2img_prep()` does the CUDA-free part (import +
+#     LoRA file check) so /config and the endpoint can report availability
+#     cheaply; `_ensure_img2img_pipe()` does the actual GPU build, lazily, the
+#     first time `_img2img_infer` runs in a worker.
+#   * @spaces.GPU pickles arguments to/from the worker. The Flux pipe is NOT
+#     picklable, so it is fetched from module state inside the worker and never
+#     crosses the boundary. Only the PIL image in and the PIL image out cross
+#     it (both picklable), exactly like `_nsfw_infer`.
+#   * This file is self-contained on the Space (no backend/ package), so the
+#     squaring + load + stylize logic from backend/hf/3_infer_img2img.py is
+#     inlined here rather than imported.
+#
+# STAGED ROLLOUT: the scene LoRA is OPTIONAL. Stage 1 (now) is "pure img2img" —
+# the base Flux2Klein model develops the photo with a generic watercolour prompt,
+# no .safetensors required. Stage 2 is dropping the LoRA file in (or setting
+# SOZAI_IMG2IMG_LORA_REPO): it's detected automatically and the trained
+# watercolour style + trigger prompt switch on. So a missing LoRA degrades to
+# base img2img rather than disabling the feature.
+#
+# Fails CLOSED (available=false) only when torch/diffusers are absent, so a CPU
+# Space or local Mac keeps the original photo instead of erroring/OOMing.
+# --------------------------------------------------------------------------- #
+IMG2IMG_MODEL_ID = os.environ.get("SOZAI_IMG2IMG_MODEL", "black-forest-labs/FLUX.2-klein-4B")
+# Path to the watercolour scene LoRA .safetensors. OPTIONAL: ship the file in the
+# repo at this path, or point SOZAI_IMG2IMG_LORA_REPO/_FILE at a Hub repo to have
+# it downloaded (CPU/network only) at prep time. Absent = base img2img (stage 1).
+IMG2IMG_LORA_PATH = os.environ.get(
+    "SOZAI_IMG2IMG_LORA",
+    os.path.join(BASE, "data/places/checkpoints/hf-spaces/scene_lora.safetensors"),
+)
+IMG2IMG_LORA_REPO = os.environ.get("SOZAI_IMG2IMG_LORA_REPO", "")
+IMG2IMG_LORA_FILE = os.environ.get("SOZAI_IMG2IMG_LORA_FILE", "scene_lora.safetensors")
+IMG2IMG_LORA_SCALE = float(os.environ.get("SOZAI_IMG2IMG_LORA_SCALE", "1.0"))
+# Generation defaults (from backend/hf/3_infer_img2img.py).
+IMG2IMG_TRIGGER = "wtrclr8"
+# Prompt used WITH the LoRA (its trigger token), vs. the base-model fallback used
+# in stage 1 when no LoRA is loaded — the trigger token means nothing without it.
+IMG2IMG_PROMPT = os.environ.get("SOZAI_IMG2IMG_PROMPT", "WTRCLR8 watercolour painting")
+IMG2IMG_BASE_PROMPT = os.environ.get(
+    "SOZAI_IMG2IMG_BASE_PROMPT",
+    "a soft watercolour painting of this scene, painterly, textured paper, gentle washes",
+)
+IMG2IMG_SIZE = int(os.environ.get("SOZAI_IMG2IMG_SIZE", "1024"))
+IMG2IMG_STEPS = int(os.environ.get("SOZAI_IMG2IMG_STEPS", "8"))
+IMG2IMG_GUIDANCE = float(os.environ.get("SOZAI_IMG2IMG_GUIDANCE", "3.5"))
+IMG2IMG_SEED = int(os.environ.get("SOZAI_IMG2IMG_SEED", "42"))
+_img2img_lock = threading.Lock()
+_img2img_state = {"loaded": False, "available": False, "pipe": None,
+                  "lora_path": None, "lora_loaded": False}
+
+
+def _img2img_prep() -> bool:
+    """CUDA-free availability check, run in the main process. Verifies diffusers
+    (with Flux2KleinPipeline) imports; the LoRA is OPTIONAL — if a file is present
+    (or downloadable via SOZAI_IMG2IMG_LORA_REPO) its path is cached for stage 2,
+    otherwise we fall back to base img2img. Does NOT touch CUDA or build the
+    pipeline. Returns True whenever the base feature can run in a GPU worker."""
+    if _img2img_state["loaded"]:
+        return _img2img_state["available"]
+    with _img2img_lock:
+        if _img2img_state["loaded"]:
+            return _img2img_state["available"]
+        try:
+            import torch  # noqa: F401  (import check only — no CUDA here)
+            from diffusers import Flux2KleinPipeline  # noqa: F401
+            # Resolve the LoRA if we can, but its absence is NOT fatal.
+            lora_path = IMG2IMG_LORA_PATH if (IMG2IMG_LORA_PATH and
+                                              os.path.isfile(IMG2IMG_LORA_PATH)) else None
+            if lora_path is None and IMG2IMG_LORA_REPO:
+                try:
+                    from huggingface_hub import hf_hub_download
+                    lora_path = hf_hub_download(repo_id=IMG2IMG_LORA_REPO,
+                                                filename=IMG2IMG_LORA_FILE)
+                except Exception as lexc:  # noqa: BLE001
+                    print(f"[img2img] LoRA download skipped ({lexc}); using base model")
+                    lora_path = None
+            _img2img_state["lora_path"] = lora_path
+            _img2img_state["available"] = True
+            print(f"[img2img] available ({IMG2IMG_MODEL_ID}; "
+                  f"LoRA={lora_path or 'none — base img2img (stage 1)'}); "
+                  f"pipeline will build on first use inside the GPU worker")
+        except Exception as exc:  # noqa: BLE001
+            print(f"[img2img] unavailable, endpoint will report disabled: {exc}")
+            _img2img_state["available"] = False
+        finally:
+            _img2img_state["loaded"] = True
+    return _img2img_state["available"]
+
+
+_img2img_build_lock = threading.Lock()
+
+
+def _ensure_img2img_pipe():
+    """Build (once) and return the Flux2KleinPipeline on the GPU, applying the
+    watercolour LoRA only if one was resolved (stage 2); otherwise it's plain
+    base img2img (stage 1). MUST be called from inside a @spaces.GPU context on
+    ZeroGPU so the CUDA context is created with a GPU attached. Idempotent +
+    thread-safe; caches into _img2img_state. Returns the pipe, or None on failure
+    (caller falls back)."""
+    pipe = _img2img_state.get("pipe")
+    if pipe is not None:
+        return pipe
+    with _img2img_build_lock:
+        pipe = _img2img_state.get("pipe")
+        if pipe is not None:
+            return pipe
+        try:
+            import torch
+            from diffusers import Flux2KleinPipeline
+            try:
+                from pillow_heif import register_heif_opener
+                register_heif_opener()  # so HEIC/HEIF phone photos open
+            except Exception:  # noqa: BLE001
+                pass
+            lora_path = _img2img_state.get("lora_path")
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            print(f"[img2img] building {IMG2IMG_MODEL_ID} on {device} "
+                  f"({'with LoRA' if lora_path else 'base, no LoRA'}) …")
+            pipe = Flux2KleinPipeline.from_pretrained(
+                IMG2IMG_MODEL_ID, torch_dtype=torch.bfloat16
+            ).to(device)
+            if lora_path and os.path.isfile(lora_path):
+                from pathlib import Path as _Path
+                lf = _Path(lora_path)
+                pipe.load_lora_weights(str(lf.parent), weight_name=lf.name,
+                                       adapter_name=IMG2IMG_TRIGGER)
+                pipe.set_adapters([IMG2IMG_TRIGGER], adapter_weights=[IMG2IMG_LORA_SCALE])
+                _img2img_state["lora_loaded"] = True
+                print("[img2img] pipeline ready (watercolour LoRA applied)")
+            else:
+                _img2img_state["lora_loaded"] = False
+                print("[img2img] pipeline ready (base model — stage 1, no LoRA)")
+            _img2img_state["pipe"] = pipe
+            return pipe
+        except Exception as exc:  # noqa: BLE001
+            print(f"[img2img] in-worker build failed: {exc}")
+            return None
+
+
+def _img2img_square(img, size: int = IMG2IMG_SIZE):
+    """EXIF-normalise, flatten to RGB, and centre-crop a PIL image to size×size —
+    the shape Klein expects (inlined from backend/hf/3_infer_img2img.py)."""
+    from PIL import Image as PILImage, ImageOps
+    src = ImageOps.exif_transpose(img).convert("RGB")
+    w, h = src.size
+    s = min(w, h)
+    src = src.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
+    return src.resize((size, size), PILImage.LANCZOS)
+
+
+@spaces.GPU(duration=120)
+def _img2img_infer(src):
+    """Develop one PIL photo into a watercolour PIL image, inside a @spaces.GPU
+    context. Builds/reuses the pipeline here (never in the main process, never
+    pickled across the boundary). `src` and the returned image are PIL (both
+    picklable), the only things that cross the GPU boundary. Generous duration
+    because the FIRST call also pays the ~14 GB cold load; later calls reuse the
+    cached pipe and are fast."""
+    import torch
+    pipe = _ensure_img2img_pipe()
+    if pipe is None:
+        raise RuntimeError("img2img pipeline could not be initialized")
+    sq = _img2img_square(src, size=IMG2IMG_SIZE)
+    # Trigger-token prompt only makes sense once the LoRA is loaded; otherwise
+    # use the generic base-model watercolour prompt (stage 1).
+    prompt = IMG2IMG_PROMPT if _img2img_state.get("lora_loaded") else IMG2IMG_BASE_PROMPT
+    device = getattr(pipe, "device", None)
+    generator = torch.Generator(
+        str(device) if device is not None else "cuda"
+    ).manual_seed(IMG2IMG_SEED)
+    result = pipe(
+        prompt=prompt,
+        image=sq,
+        height=IMG2IMG_SIZE,
+        width=IMG2IMG_SIZE,
+        num_inference_steps=IMG2IMG_STEPS,
+        guidance_scale=IMG2IMG_GUIDANCE,
+        generator=generator,
+    )
+    return result.images[0]
+
+
+def _pil_to_data_url(img, fmt: str = "PNG") -> str:
+    import base64
+    buf = io.BytesIO()
+    img.save(buf, format=fmt)
+    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+    return f"data:image/{fmt.lower()};base64,{b64}"
+
+
+@app.post("/api/img2img")
+async def img2img_rest(request: Request):
+    """Develop one uploaded photo into a watercolour painting.
+    Body: {data_url: "data:image/...;base64,...", name?: str}.
+    Returns {available, image: <png data url>, model}. When the pipeline can't
+    run on this host it returns {available: false, image: null} so the caller can
+    keep the original photo rather than treating it as an error."""
+    import base64
+    if not _img2img_prep():
+        return {"available": False, "image": None, "model": IMG2IMG_MODEL_ID}
+    try:
+        body = await request.json()
+    except Exception:  # noqa: BLE001
+        raise HTTPException(status_code=400, detail="invalid JSON body")
+    data_url = (body or {}).get("data_url", "")
+    name = str((body or {}).get("name", "upload"))[:60]
+    if not data_url:
+        raise HTTPException(status_code=400, detail="missing data_url")
+    try:
+        b64 = data_url.split(",", 1)[1] if "," in data_url else data_url
+        raw = base64.b64decode(b64)
+        from PIL import Image
+        src = Image.open(io.BytesIO(raw))
+        with traced("img2img", model=IMG2IMG_MODEL_ID, file=name) as t:
+            out = _img2img_infer(src)
+            t.set_output({"size": list(out.size)})
+        return {"available": True, "image": _pil_to_data_url(out), "model": IMG2IMG_MODEL_ID}
+    except HTTPException:
+        raise
+    except Exception as exc:  # noqa: BLE001
+        print(f"[img2img] generation failed: {exc}")
+        raise HTTPException(status_code=500, detail=str(exc))
+
+
 # --------------------------------------------------------------------------- #
 # Voice → text via NVIDIA NeMo ASR (item 5). Streaming Nemotron model; needs a
 # GPU to be practical. Loaded lazily; if NeMo/torch/GPU are missing the endpoint

packages.txt

@@ -0,0 +1 @@
+libheif-dev

requirements.txt

@@ -1,8 +1,11 @@
 gradio
 uvicorn[standard]
 transformers
-diffusers
+# Flux2KleinPipeline (img2img watercolour develop) is not in a released diffusers
+# yet — it must come from git main. See backend/hf/3_infer_img2img.py.
+diffusers @ git+https://github.com/huggingface/diffusers.git
 accelerate
+pillow-heif   # HEIC/HEIF source photos for the img2img develop step
 hf_xet
 huggingface_hub
 soundfile>=0.12