Refactor code structure for improved readability and maintainability

runpod-worker/Dockerfile

@@ -0,0 +1,29 @@
+# RunPod Serverless worker: FLUX.2 [klein] image editing.
+# CUDA 12.4 base + torch cu124: RunPod's RTX 3090/4090 hosts run driver 550
+# (CUDA 12.4). A 12.8 base demands driver >= 12.8 and the container REFUSES to
+# start on a 550 host ("nvidia-container-cli: unsatisfied condition: cuda>=12.8").
+# 12.4 runs on driver 550+ and carries sm_86 (3090) + sm_89 (4090) kernels.
+# NOTE: this does NOT carry Blackwell (sm_120) kernels — keep this endpoint's GPU
+# pool to 3090/4090 (24 GB Ada/Ampere), not the PRO 6000 Blackwell cards.
+FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        python3 python3-pip git && \
+    ln -sf /usr/bin/python3 /usr/bin/python && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# torch (CUDA 12.4 build) FIRST — matches the 12.4 base + driver 550 hosts.
+RUN pip install --upgrade pip && \
+    pip install torch --index-url https://download.pytorch.org/whl/cu124
+COPY runpod-worker/requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY runpod-worker/handler.py .
+
+CMD ["python3", "-u", "handler.py"]

runpod-worker/README.md

@@ -0,0 +1,49 @@
+# Instruct-Pix2Pix image-editing worker (RunPod Serverless)
+
+Instruction-based image editing that powers the gallery editor's **Apply Prompt**
+button. Instruct-Pix2Pix is *trained to follow edit instructions* ("make it
+golden hour", "add snow", "turn the sky orange") and applies them while keeping
+the rest of the photo — unlike plain SD img2img.
+
+Small (~2 GB), runs on your existing 24 GB endpoint, **no gated license, no HF
+token, no volume resize**. Input/output matches the app, so no app-code change.
+
+## Contract
+
+Request → `POST /runsync`:
+```json
+{ "input": { "image": "<base64>", "prompt": "<edit instruction>",
+             "guidance_scale": 7.5, "num_inference_steps": 25, "seed": 123 } }
+```
+Response → `{ "output": { "image": "<base64 PNG>" } }` (or `{ "output": { "error": "..." } }`).
+
+`strength` / `mask` from the app are ignored (this model conditions on the image
+via `image_guidance_scale`, not img2img noise or masks).
+
+## Deploy (reuse your existing endpoint)
+
+No RunPod settings to change — just push, and RunPod auto-rebuilds:
+```bash
+cd "advance-photo-gallery-web-sdk - Copy"
+git add runpod-worker
+git commit -m "Switch worker to Instruct-Pix2Pix"
+git push
+```
+Then **warm up once** (downloads ~2 GB — quick):
+```bash
+curl -X POST https://api.runpod.ai/v2/<ENDPOINT_ID>/runsync \
+  -H "Content-Type: application/json" -H "Authorization: Bearer <API_KEY>" \
+  -d @runpod-worker/warmup.json
+```
+Test **Apply Prompt** in the app — no Vercel change (same URL).
+
+## Tuning (optional endpoint env vars)
+
+| Var | Default | Purpose |
+|---|---|---|
+| `IP2P_MODEL` | `timbrooks/instruct-pix2pix` | the model |
+| `IP2P_IMAGE_GUIDANCE` | `1.5` | higher keeps more of the original; **lower toward ~1.2 if edits feel too weak** |
+| `IP2P_MAX_SIZE` | `768` | max long-edge px (SD1.5 works best ≤ 768) |
+
+**Prompts work best as imperatives:** "make it a golden-hour sunset", "add snow",
+"turn it into winter" — not descriptions.

runpod-worker/handler.py

@@ -0,0 +1,113 @@
+"""
+RunPod Serverless worker: FLUX.2 [klein] 4B instruction image editing (img2img).
+
+Replaces the old Instruct-Pix2Pix worker. FLUX.2-klein is Apache-2.0 (no HF token),
+runs on a 24 GB card (RTX 3090), and is distilled → ~6 steps, fast. It needs a
+recent torch, which also carries kernels for the newest (Blackwell) GPUs — so it
+runs on whatever card RunPod assigns.
+
+Matches the app's img2img contract (no app-code change needed):
+  {"input": {
+     "image": "<base64>",   # required (raw base64 or data: URI)
+     "prompt": "<edit instruction>",  # required
+     "strength": 0.6,       # optional — the editor's "Edit strength" slider (mapped to guidance)
+     "seed": 123            # optional
+  }}
+Returns: {"image": "<base64 PNG>"}  (or {"error": "..."}).
+
+Tuning (env): FLUX2_MODEL (default black-forest-labs/FLUX.2-klein-4B),
+FLUX2_STEPS (default 6), FLUX2_MAX_SIZE (default 1024).
+"""
+
+import base64
+import io
+import os
+
+_V = "/runpod-volume"
+if os.path.isdir(_V):
+    os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
+
+import torch  # noqa: E402
+import runpod  # noqa: E402
+from PIL import Image  # noqa: E402
+from diffusers import Flux2KleinPipeline  # noqa: E402
+
+MODEL = os.environ.get("FLUX2_MODEL", "black-forest-labs/FLUX.2-klein-4B")
+STEPS = int(os.environ.get("FLUX2_STEPS", "6"))
+MAX_SIZE = int(os.environ.get("FLUX2_MAX_SIZE", "1024"))
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+_pipe = None
+
+
+def _get_pipe():
+    global _pipe
+    if _pipe is None:
+        pipe = Flux2KleinPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
+        total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if DEVICE == "cuda" else 0
+        if total_gb >= 20:
+            pipe = pipe.to("cuda")
+        else:
+            pipe.enable_model_cpu_offload()
+        pipe.set_progress_bar_config(disable=True)
+        _pipe = pipe
+    return _pipe
+
+
+def _b64_to_image(data):
+    if "," in data:
+        data = data.split(",", 1)[1]
+    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")
+
+
+def _image_to_b64(img):
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+
+def _fit(img):
+    w, h = img.size
+    scale = min(1.0, MAX_SIZE / max(w, h))
+    nw = max(16, (int(w * scale) // 16) * 16)
+    nh = max(16, (int(h * scale) // 16) * 16)
+    return img.resize((nw, nh), Image.LANCZOS)
+
+
+def handler(event):
+    inp = (event or {}).get("input") or {}
+    image_b64 = inp.get("image")
+    prompt = (inp.get("prompt") or "").strip()
+    if not image_b64:
+        return {"error": "Missing 'image' (base64)."}
+    if not prompt:
+        return {"error": "Missing 'prompt' (edit instruction)."}
+
+    try:
+        image = _fit(_b64_to_image(image_b64))
+        # The editor's "Edit strength" slider (0..1) maps to guidance — FLUX.2 klein
+        # likes low guidance (~1); stronger = higher.
+        guidance = 1.0
+        strength = inp.get("strength")
+        if strength is not None:
+            s = max(0.0, min(1.0, float(strength)))
+            guidance = round(1.0 + s * 3.0, 2)
+        seed = inp.get("seed")
+        generator = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else None
+
+        result = _get_pipe()(
+            prompt=prompt,
+            image=image,
+            height=image.height,
+            width=image.width,
+            guidance_scale=guidance,
+            num_inference_steps=STEPS,
+            generator=generator,
+        ).images[0]
+
+        return {"image": _image_to_b64(result)}
+    except Exception as exc:
+        return {"error": f"{type(exc).__name__}: {exc}"}
+
+
+runpod.serverless.start({"handler": handler})

runpod-worker/requirements.txt

@@ -0,0 +1,12 @@
+# FLUX.2 [klein] worker. torch is installed (CUDA 12.8 build) FIRST in the
+# Dockerfile so it carries Blackwell kernels; diffusers is from git for the
+# Flux2KleinPipeline. FLUX.2 uses a Mistral-based text encoder → transformers +
+# sentencepiece/protobuf.
+diffusers @ git+https://github.com/huggingface/diffusers.git
+transformers>=4.49.0
+accelerate>=1.2.0
+safetensors>=0.4.5
+sentencepiece>=0.2.0
+protobuf>=4.25.0
+Pillow==10.4.0
+runpod==1.7.9

runpod-worker/test_input.json

@@ -0,0 +1,8 @@
+{
+  "input": {
+    "image": "PASTE_BASE64_IMAGE_HERE",
+    "prompt": "a vibrant golden-hour sky, dramatic clouds, photorealistic",
+    "strength": 0.6,
+    "num_inference_steps": 25
+  }
+}

runpod-worker/warmup.json

@@ -0,0 +1 @@
+{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg==", "prompt": "turn the sky into a dramatic golden-hour sunset with warm orange clouds"}}