"""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: # noqa: BLE001 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: # noqa: BLE001 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: # noqa: BLE001 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: # noqa: BLE001 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, }, )