| """Generic gated Replicate runner + FLUX Kontext beautify (instruction img2img). |
| |
| FLUX Kontext is a strong instruction image editor (far stronger than |
| gemini-2.5-flash-image at "edit the face, keep the rest"). It re-renders the |
| frame, so hair/background are visually preserved (prompt-named) but not byte- |
| exact; we upscale + GFPGAN afterwards. Gated like every paid path: a real call |
| happens only when REPLICATE_API_TOKEN is present (the wrapper sets it for the run |
| and clears it after). The token is never logged, stored, or committed. |
| |
| Config (env, never committed): |
| REPLICATE_API_TOKEN=... # this shell only |
| KONTEXT_MODEL=black-forest-labs/flux-kontext-pro # override e.g. -max |
| Pilot Ready: NOT CONFIRMED. |
| """ |
| from __future__ import annotations |
|
|
| import base64 |
| import os |
| import time |
|
|
| REPLICATE_API = "https://api.replicate.com/v1" |
|
|
|
|
| class ReplicateUnavailable(RuntimeError): |
| """A Replicate call is disabled (no token), misconfigured, or failed.""" |
|
|
|
|
| def replicate_token() -> str: |
| return (os.getenv("REPLICATE_API_TOKEN") or "").strip() |
|
|
|
|
| def replicate_real_enabled() -> bool: |
| """A real Replicate call is only made when a token is present.""" |
| return bool(replicate_token()) |
|
|
|
|
| def _data_uri(image_bytes: bytes, mime: str = "image/png") -> str: |
| return f"data:{mime};base64," + base64.b64encode(image_bytes).decode("ascii") |
|
|
|
|
| def _safe_err(text: str, limit: int = 200) -> str: |
| return (text or "").strip().replace("\n", " ")[:limit] |
|
|
|
|
| def run_replicate_model( |
| model: str, |
| model_input: dict, |
| *, |
| version: str | None = None, |
| timeout_s: int = 300, |
| poll_s: float = 2.0, |
| timings: dict | None = None, |
| ) -> bytes: |
| """Run any Replicate model and return the first output image bytes. |
| |
| Resolves the model's latest version (free GET) unless one is given, creates a |
| prediction, polls to completion, and downloads the output image. Gated on the |
| token; raises ReplicateUnavailable on any problem. Token never logged. |
| |
| If `timings` (a dict) is passed, the hosted round-trip is split into |
| `fal_submit` / `fal_inference` / `download` seconds (reusing the engine's stage |
| names so the latency telemetry is identical across backends). Optional and |
| backward-compatible — callers that omit it are unaffected. |
| """ |
| if not replicate_real_enabled(): |
| raise ReplicateUnavailable( |
| "REPLICATE_API_TOKEN not set; no network call made" |
| ) |
| token = replicate_token() |
| try: |
| import httpx |
| except Exception as exc: |
| raise ReplicateUnavailable("httpx is not installed") from exc |
|
|
| headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} |
| ver = (version or os.getenv("REPLICATE_MODEL_VERSION") or "").strip() |
|
|
| try: |
| with httpx.Client(timeout=30.0) as client: |
| if not ver: |
| mr = client.get(f"{REPLICATE_API}/models/{model}", headers=headers) |
| if mr.status_code == 404: |
| raise ReplicateUnavailable( |
| f"model '{model}' not found on Replicate (check the slug)" |
| ) |
| if mr.status_code >= 400: |
| raise ReplicateUnavailable( |
| f"model lookup failed: HTTP {mr.status_code} {_safe_err(mr.text)}" |
| ) |
| ver = ((mr.json() or {}).get("latest_version") or {}).get("id") or "" |
| if not ver: |
| raise ReplicateUnavailable(f"model '{model}' has no usable version") |
|
|
| t_submit = time.time() |
| r = client.post(f"{REPLICATE_API}/predictions", headers=headers, |
| json={"version": ver, "input": model_input}) |
| if r.status_code >= 400: |
| raise ReplicateUnavailable( |
| f"create failed: HTTP {r.status_code} {_safe_err(r.text)}" |
| ) |
| pred = r.json() |
| get_url = (pred.get("urls") or {}).get("get") |
| status = pred.get("status") |
| if timings is not None: |
| timings["fal_submit"] = round(time.time() - t_submit, 3) |
|
|
| t_inf = time.time() |
| waited = 0.0 |
| while status not in {"succeeded", "failed", "canceled"} and get_url: |
| if waited >= timeout_s: |
| raise ReplicateUnavailable(f"prediction timed out after {timeout_s}s") |
| time.sleep(poll_s) |
| waited += poll_s |
| pred = client.get(get_url, headers=headers).json() |
| status = pred.get("status") |
| if timings is not None: |
| timings["fal_inference"] = round(time.time() - t_inf, 3) |
|
|
| if status != "succeeded": |
| raise ReplicateUnavailable( |
| f"prediction {status}: {_safe_err(str(pred.get('error')))}" |
| ) |
|
|
| out = pred.get("output") |
| url = out[-1] if isinstance(out, list) and out else (out if isinstance(out, str) else None) |
| if not url: |
| raise ReplicateUnavailable("prediction returned no image URL") |
| t_dl = time.time() |
| img = client.get(url, timeout=60.0) |
| if img.status_code >= 400 or not img.content: |
| raise ReplicateUnavailable(f"output download failed: HTTP {img.status_code}") |
| if timings is not None: |
| timings["download"] = round(time.time() - t_dl, 3) |
| return img.content |
| except ReplicateUnavailable: |
| raise |
| except Exception as exc: |
| raise ReplicateUnavailable(f"Replicate call failed: {type(exc).__name__}") from exc |
|
|
|
|
| def fetch_model_schema(model: str) -> dict: |
| """Return the model's input field properties (free GET; no prediction billed). |
| |
| Lets a runner auto-match field names (image / prompt / strength) instead of |
| guessing and hitting a 422. Gated on the token; raises on problems. |
| """ |
| if not replicate_real_enabled(): |
| raise ReplicateUnavailable("REPLICATE_API_TOKEN not set") |
| token = replicate_token() |
| try: |
| import httpx |
| except Exception as exc: |
| raise ReplicateUnavailable("httpx is not installed") from exc |
| try: |
| with httpx.Client(timeout=30.0) as client: |
| r = client.get(f"{REPLICATE_API}/models/{model}", |
| headers={"Authorization": f"Bearer {token}"}) |
| if r.status_code >= 400: |
| raise ReplicateUnavailable( |
| f"model lookup failed: HTTP {r.status_code} {_safe_err(r.text)}" |
| ) |
| ver = (r.json() or {}).get("latest_version") or {} |
| return (((ver.get("openapi_schema") or {}).get("components") or {}) |
| .get("schemas", {}).get("Input", {}).get("properties", {})) or {} |
| except ReplicateUnavailable: |
| raise |
| except Exception as exc: |
| raise ReplicateUnavailable(f"schema fetch failed: {type(exc).__name__}") from exc |
|
|
|
|
| def replicate_fill_model() -> str: |
| return (os.getenv("REPLICATE_FILL_MODEL") or "black-forest-labs/flux-fill-dev").strip() |
|
|
|
|
| def run_replicate_inpaint( |
| image_bytes: bytes, |
| mask_bytes: bytes, |
| prompt: str, |
| *, |
| num_inference_steps: int = 28, |
| guidance: float = 30.0, |
| timings: dict | None = None, |
| ) -> bytes: |
| """Masked FLUX inpaint on Replicate (flux-fill-dev): regenerate the WHITE mask |
| region only, keep everything else — so hair / clothing / background outside the |
| face mask are preserved, exactly like the fal path. Returns the result bytes. |
| |
| `mask_bytes` is an L/grayscale PNG (white = regenerate). Field names follow the |
| flux-fill schema (image / mask / prompt); override via REPLICATE_FILL_*_FIELD if |
| a fork differs. Gated on REPLICATE_API_TOKEN; the token is never logged. |
| |
| NOTE: flux-fill has NO IP-Adapter and NO strength dial — it regenerates the |
| masked face from the prompt + surrounding context, so identity preservation is |
| weaker/more aggressive than fal flux-general. Verify on a real face before use. |
| """ |
| model_input = { |
| os.getenv("REPLICATE_FILL_IMAGE_FIELD", "image"): _data_uri(image_bytes), |
| os.getenv("REPLICATE_FILL_MASK_FIELD", "mask"): _data_uri(mask_bytes), |
| "prompt": prompt, |
| "num_inference_steps": int(num_inference_steps), |
| "guidance": float(guidance), |
| "output_format": "png", |
| } |
| return run_replicate_model(replicate_fill_model(), model_input, timings=timings) |
|
|
|
|
| def flux_kontext_model() -> str: |
| return (os.getenv("KONTEXT_MODEL") or "black-forest-labs/flux-kontext-pro").strip() |
|
|
|
|
| def flux_kontext_edit(image_bytes: bytes, prompt: str) -> bytes: |
| """Instruction-edit `image_bytes` with FLUX Kontext and return the result.""" |
| return run_replicate_model( |
| flux_kontext_model(), |
| { |
| "input_image": _data_uri(image_bytes), |
| "prompt": prompt, |
| "aspect_ratio": "match_input_image", |
| "output_format": "png", |
| "safety_tolerance": 2, |
| }, |
| ) |
|
|
