mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-08-02 11:38:40 -04:00
pewdiepie-archdaemon
1507 lines
61 KiB
Python
1507 lines
61 KiB
Python
#!/usr/bin/env python3
|
||
"""Minimal OpenAI-compatible image generation API server using diffusers.
|
||
|
||
Serves /v1/images/generations and /v1/models for compatibility with
|
||
Odysseus's image generation tool.
|
||
|
||
Usage:
|
||
python3 scripts/diffusion_server.py --model /path/to/model --port 8100
|
||
"""
|
||
import os
|
||
import sys
|
||
import importlib
|
||
import importlib.machinery
|
||
# Block xformers — create a fake module that reports as not installed
|
||
_fake = type(sys)("xformers")
|
||
_fake.__version__ = "0.0.0"
|
||
_fake.__spec__ = importlib.machinery.ModuleSpec("xformers", None)
|
||
_fake.__path__ = []
|
||
sys.modules["xformers"] = _fake
|
||
sys.modules["xformers.ops"] = type(sys)("xformers.ops")
|
||
sys.modules["xformers.ops.fmha"] = type(sys)("xformers.ops.fmha")
|
||
|
||
import argparse
|
||
import base64
|
||
import io
|
||
import json
|
||
import logging
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
from contextlib import asynccontextmanager
|
||
|
||
import torch
|
||
import uvicorn
|
||
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from starlette.middleware.trustedhost import TrustedHostMiddleware
|
||
from pydantic import BaseModel
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger("diffusion_server")
|
||
|
||
_pipe = None
|
||
_model_id = ""
|
||
DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
|
||
_args = None
|
||
_PROGRESS = {}
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(application):
|
||
load_model()
|
||
yield
|
||
|
||
|
||
app = FastAPI(title="Diffusion Server", lifespan=lifespan)
|
||
|
||
# Conservative defaults — server is designed for server-to-server use from
|
||
# the Odysseus backend. Wildcard CORS + the 127.0.0.1 default bind used to
|
||
# leave the server reachable via DNS-rebinding from any browser tab on the
|
||
# same host. The CLI flags below extend these allowlists for operators who
|
||
# need browser access; the safe defaults handle the common case.
|
||
_DEFAULT_ALLOWED_HOSTS = ["127.0.0.1", "localhost", "::1"]
|
||
_DEFAULT_CORS_ORIGINS: list = [] # default-deny
|
||
|
||
|
||
def _compute_allowed_hosts(bind_host: str, extras=None) -> list:
|
||
"""Allowed Host header values: the bind address + loopback variants +
|
||
any operator-supplied --allowed-host values. Duplicates and empty
|
||
strings are dropped; order is stable for predictable middleware setup."""
|
||
seen = []
|
||
for h in (bind_host, *_DEFAULT_ALLOWED_HOSTS, *(extras or [])):
|
||
h = (h or "").strip()
|
||
if h and h not in seen:
|
||
seen.append(h)
|
||
return seen
|
||
|
||
|
||
def _compute_cors_origins(extras=None) -> list:
|
||
"""CORS allowlist: default-deny (empty), extended only by explicit
|
||
--allowed-origin values. Server-to-server callers don't set an Origin
|
||
header so they're unaffected; this only narrows browser access."""
|
||
seen = []
|
||
for o in (*_DEFAULT_CORS_ORIGINS, *(extras or [])):
|
||
o = (o or "").strip()
|
||
if o and o not in seen:
|
||
seen.append(o)
|
||
return seen
|
||
|
||
|
||
def _configure_security_middleware(application, allowed_hosts, allowed_origins):
|
||
"""Replace `application`'s user middleware stack with the diffusion server
|
||
security middleware: the TrustedHost allowlist and, when origins are
|
||
supplied, CORS. Used at module load and by the __main__ CLI path before
|
||
serving starts. Raises before mutating if the middleware stack has already
|
||
been built. Order is preserved: TrustedHost first, then CORS (added last ->
|
||
outermost)."""
|
||
if application.middleware_stack is not None:
|
||
raise RuntimeError("security middleware must be configured before the app starts serving")
|
||
application.user_middleware.clear()
|
||
application.add_middleware(TrustedHostMiddleware, allowed_hosts=list(allowed_hosts))
|
||
if allowed_origins:
|
||
application.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=list(allowed_origins),
|
||
allow_methods=["GET", "POST", "OPTIONS"],
|
||
allow_headers=["Authorization", "Content-Type"],
|
||
)
|
||
|
||
|
||
# Install defaults at module load so importing the app for tests / direct
|
||
# uvicorn invocation still benefits from the Host-header allowlist.
|
||
_configure_security_middleware(app, _DEFAULT_ALLOWED_HOSTS, _DEFAULT_CORS_ORIGINS)
|
||
|
||
|
||
def _start_progress(request_id: str, total_steps: int, prompt: str, kind: str) -> str:
|
||
rid = (request_id or "").strip() or uuid.uuid4().hex
|
||
_PROGRESS[rid] = {
|
||
"id": rid,
|
||
"kind": kind,
|
||
"status": "running",
|
||
"step": 0,
|
||
"total": max(1, int(total_steps or 1)),
|
||
"percent": 0,
|
||
"prompt": (prompt or "")[:120],
|
||
"started_at": time.time(),
|
||
"updated_at": time.time(),
|
||
}
|
||
return rid
|
||
|
||
|
||
def _update_progress(request_id: str, step: int, total_steps: int | None = None, status: str = "running"):
|
||
if not request_id:
|
||
return
|
||
item = _PROGRESS.get(request_id)
|
||
if not item:
|
||
return
|
||
total = max(1, int(total_steps or item.get("total") or 1))
|
||
current = max(0, min(int(step or 0), total))
|
||
item.update({
|
||
"status": status,
|
||
"step": current,
|
||
"total": total,
|
||
"percent": round((current / total) * 100, 1),
|
||
"updated_at": time.time(),
|
||
})
|
||
|
||
|
||
def _finish_progress(request_id: str, status: str = "done", error: str = ""):
|
||
if not request_id:
|
||
return
|
||
item = _PROGRESS.get(request_id)
|
||
if not item:
|
||
return
|
||
total = max(1, int(item.get("total") or 1))
|
||
item.update({
|
||
"status": status,
|
||
"step": total if status == "done" else item.get("step", 0),
|
||
"total": total,
|
||
"percent": 100 if status == "done" else item.get("percent", 0),
|
||
"error": error,
|
||
"updated_at": time.time(),
|
||
})
|
||
|
||
|
||
def _run_pipeline_with_progress(pipe, request_id: str, total_steps: int, **kwargs):
|
||
def step_end_callback(_pipe, step, timestep, callback_kwargs):
|
||
_update_progress(request_id, int(step) + 1, total_steps)
|
||
return callback_kwargs
|
||
|
||
def legacy_callback(step, timestep, latents):
|
||
_update_progress(request_id, int(step) + 1, total_steps)
|
||
|
||
try:
|
||
return pipe(callback_on_step_end=step_end_callback, **kwargs)
|
||
except TypeError as exc:
|
||
if "callback_on_step_end" not in str(exc):
|
||
raise
|
||
try:
|
||
return pipe(callback=legacy_callback, callback_steps=1, **kwargs)
|
||
except TypeError as exc:
|
||
if "callback" not in str(exc) and "callback_steps" not in str(exc):
|
||
raise
|
||
return pipe(**kwargs)
|
||
|
||
|
||
class ImageRequest(BaseModel):
|
||
model: str = ""
|
||
prompt: str
|
||
n: int = 1
|
||
size: str = "1024x1024"
|
||
quality: str = "medium"
|
||
response_format: str = "b64_json"
|
||
request_id: str = ""
|
||
|
||
|
||
def _parse_size(size: str) -> tuple[int, int]:
|
||
try:
|
||
w, h = (size or "").split("x")
|
||
return int(w), int(h)
|
||
except Exception:
|
||
return _args.width, _args.height
|
||
|
||
|
||
def _quality_steps(quality: str) -> int:
|
||
default_steps = _args.steps or 8
|
||
steps_map = {"low": 4, "medium": default_steps, "high": 20, "auto": 12}
|
||
return steps_map.get(quality, default_steps)
|
||
|
||
|
||
def _guidance_scale() -> float:
|
||
return float(_args.guidance_scale)
|
||
|
||
|
||
def _default_negative_prompt() -> str | None:
|
||
value = (_args.negative_prompt or "").strip()
|
||
return value or None
|
||
|
||
|
||
def _pipeline_accepts_arg(name: str) -> bool:
|
||
try:
|
||
import inspect
|
||
sig = inspect.signature(_pipe.__call__)
|
||
return name in sig.parameters
|
||
except Exception:
|
||
return True
|
||
|
||
|
||
def _pipeline_call_kwargs(**kwargs) -> dict:
|
||
"""Filter kwargs to the active pipeline signature.
|
||
|
||
Diffusers/community pipelines are not consistent: ordinary img2img
|
||
usually accepts `image`, while instruction-edit models such as OmniGen2
|
||
accept `input_images` plus model-specific guidance fields. Filtering keeps
|
||
the server generic and avoids hardcoding repo IDs.
|
||
"""
|
||
try:
|
||
import inspect
|
||
sig = inspect.signature(_pipe.__call__)
|
||
params = sig.parameters
|
||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()):
|
||
return {k: v for k, v in kwargs.items() if v is not None}
|
||
return {k: v for k, v in kwargs.items() if v is not None and k in params}
|
||
except Exception:
|
||
return {k: v for k, v in kwargs.items() if v is not None}
|
||
|
||
|
||
def _image_response(images) -> dict:
|
||
data = []
|
||
for img in images:
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})
|
||
return {"created": int(time.time()), "data": data}
|
||
|
||
|
||
def _fix_meta_tensors(pipe, dtype):
|
||
"""Replace any meta tensors with real zero tensors on CPU so .to(device) works."""
|
||
for name, component in pipe.components.items():
|
||
if not hasattr(component, 'parameters'):
|
||
continue
|
||
fixed = 0
|
||
for pname, param in component.named_parameters():
|
||
if param.device.type == 'meta':
|
||
with torch.no_grad():
|
||
new_param = torch.zeros(param.shape, dtype=dtype, device='cpu')
|
||
# Walk to the actual module holding this param
|
||
parts = pname.split('.')
|
||
mod = component
|
||
for p in parts[:-1]:
|
||
mod = getattr(mod, p)
|
||
setattr(mod, parts[-1], torch.nn.Parameter(new_param, requires_grad=param.requires_grad))
|
||
fixed += 1
|
||
if fixed:
|
||
logger.info(f" Fixed {fixed} meta tensors in {name}")
|
||
|
||
|
||
def _target_device() -> str:
|
||
"""Best available torch device for Diffusers on this host."""
|
||
try:
|
||
if torch.cuda.is_available():
|
||
return "cuda"
|
||
except Exception:
|
||
pass
|
||
try:
|
||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||
return "mps"
|
||
except Exception:
|
||
pass
|
||
return "cpu"
|
||
|
||
|
||
def _can_cpu_offload(device: str) -> bool:
|
||
# Diffusers CPU offload helpers are CUDA/accelerate-oriented. On Apple
|
||
# Silicon they either no-op poorly or fail; MPS should use direct .to("mps").
|
||
return device == "cuda"
|
||
|
||
|
||
def _load_omnigen2_pipeline(model_path: str, torch_dtype, target_device: str, use_offload: bool) -> bool:
|
||
"""Load OmniGen2 from the official repo package.
|
||
|
||
OmniGen2 publishes a Diffusers-style model_index.json, but current
|
||
public diffusers builds do not expose OmniGen2Pipeline as a stock class.
|
||
The official examples import the pipeline from the cloned `omnigen2`
|
||
package, so support that path without hardcoding any private model.
|
||
"""
|
||
global _pipe
|
||
try:
|
||
from omnigen2.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline
|
||
from omnigen2.models.transformers.transformer_omnigen2 import OmniGen2Transformer2DModel
|
||
except Exception as exc:
|
||
logger.warning("OmniGen2 package import failed: %s", exc)
|
||
return False
|
||
|
||
try:
|
||
logger.info("Loading OmniGen2 pipeline via official omnigen2 package")
|
||
pipe = OmniGen2Pipeline.from_pretrained(
|
||
model_path,
|
||
torch_dtype=torch_dtype,
|
||
trust_remote_code=True,
|
||
)
|
||
pipe.transformer = OmniGen2Transformer2DModel.from_pretrained(
|
||
model_path,
|
||
subfolder="transformer",
|
||
torch_dtype=torch_dtype,
|
||
)
|
||
if use_offload and _can_cpu_offload(target_device):
|
||
pipe.enable_model_cpu_offload()
|
||
logger.info("Loaded OmniGen2 with CPU offload")
|
||
else:
|
||
pipe = pipe.to(target_device)
|
||
logger.info("Loaded OmniGen2 on %s", target_device)
|
||
_pipe = pipe
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning("Official OmniGen2 loader failed: %s", exc)
|
||
_pipe = None
|
||
return False
|
||
|
||
|
||
def load_model():
|
||
global _pipe, _model_id
|
||
import diffusers
|
||
|
||
model_path = _args.model
|
||
_model_id = Path(model_path).name
|
||
dtype_map = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
|
||
torch_dtype = dtype_map.get(_args.dtype, torch.bfloat16)
|
||
use_offload = _args.cpu_offload
|
||
target_device = _target_device()
|
||
if target_device != "cuda" and use_offload:
|
||
logger.warning("CPU offload requested but %s is the active device; using direct device placement instead", target_device)
|
||
use_offload = False
|
||
|
||
logger.info(f"Loading model from {model_path} (dtype={_args.dtype}, offload={use_offload}, device={target_device})...")
|
||
|
||
# Ensure HF token is available for gated repos
|
||
_hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
||
if _hf_token:
|
||
logger.info("HF token found in environment")
|
||
# Login so all huggingface_hub calls use the token
|
||
try:
|
||
from huggingface_hub import login
|
||
login(token=_hf_token, add_to_git_credential=False)
|
||
logger.info("Logged in to HuggingFace Hub")
|
||
except Exception as e:
|
||
logger.warning(f"HF login failed: {e}")
|
||
else:
|
||
logger.warning("No HF_TOKEN set — gated models will fail")
|
||
|
||
# Detect pipeline class from model_index.json
|
||
model_index = Path(model_path) / "model_index.json"
|
||
pipeline_cls = None
|
||
cls_name_from_index = ""
|
||
if model_index.exists():
|
||
try:
|
||
idx = json.loads(model_index.read_text(encoding="utf-8"))
|
||
cls_name_from_index = idx.get("_class_name", "")
|
||
if hasattr(diffusers, cls_name_from_index):
|
||
pipeline_cls = getattr(diffusers, cls_name_from_index)
|
||
logger.info(f"Detected pipeline class: {cls_name_from_index}")
|
||
else:
|
||
logger.warning(f"model_index.json says {cls_name_from_index} but not in diffusers")
|
||
except Exception as e:
|
||
logger.warning(f"Could not parse model_index.json: {e}")
|
||
|
||
# Build candidate list: detected class first, then DiffusionPipeline (auto-detect from model_index.json)
|
||
# Only try Flux-specific pipelines if model name suggests Flux
|
||
candidates = []
|
||
if pipeline_cls:
|
||
candidates.append((pipeline_cls, pipeline_cls.__name__))
|
||
# DiffusionPipeline reads model_index.json and auto-selects the right pipeline
|
||
candidates.append((diffusers.DiffusionPipeline, "DiffusionPipeline"))
|
||
# Flux-specific fallbacks only if model name hints at Flux
|
||
_model_lower = Path(model_path).name.lower()
|
||
if "flux" in _model_lower:
|
||
for name in ("Flux2Pipeline", "FluxPipeline"):
|
||
cls = getattr(diffusers, name, None)
|
||
if cls and cls not in [c for c, _ in candidates]:
|
||
candidates.append((cls, name))
|
||
|
||
def _cleanup():
|
||
import gc; gc.collect()
|
||
try:
|
||
torch.cuda.empty_cache()
|
||
logger.debug("GPU cache cleared")
|
||
except Exception as e:
|
||
logger.debug(f"GPU cache clear failed: {e}")
|
||
|
||
loaded = False
|
||
if cls_name_from_index == "OmniGen2Pipeline" or "omnigen2" in str(model_path).lower():
|
||
loaded = _load_omnigen2_pipeline(model_path, torch_dtype, target_device, use_offload)
|
||
|
||
def _load_pipe(cls, name):
|
||
"""Try loading pipeline, handling meta tensor issues."""
|
||
global _pipe
|
||
|
||
# First try normal load
|
||
try:
|
||
kwargs = {"torch_dtype": torch_dtype}
|
||
if name == "DiffusionPipeline" and cls_name_from_index and not hasattr(diffusers, cls_name_from_index):
|
||
kwargs["trust_remote_code"] = True
|
||
_pipe = cls.from_pretrained(model_path, **kwargs)
|
||
except Exception as e:
|
||
logger.warning(f"{name} from_pretrained failed: {e}")
|
||
if name == "DiffusionPipeline" and cls_name_from_index and not hasattr(diffusers, cls_name_from_index):
|
||
try:
|
||
logger.info(f"Retrying {name} with custom_pipeline={model_path}")
|
||
_pipe = cls.from_pretrained(
|
||
model_path,
|
||
torch_dtype=torch_dtype,
|
||
custom_pipeline=model_path,
|
||
trust_remote_code=True,
|
||
)
|
||
except Exception as e2:
|
||
logger.warning(f"{name} custom_pipeline retry failed: {e2}")
|
||
_pipe = None
|
||
_cleanup()
|
||
return False
|
||
else:
|
||
_pipe = None
|
||
_cleanup()
|
||
return False
|
||
|
||
# Materialize any meta tensors before moving to device
|
||
_fix_meta_tensors(_pipe, torch_dtype)
|
||
|
||
if use_offload and _can_cpu_offload(target_device):
|
||
try:
|
||
_pipe.enable_model_cpu_offload()
|
||
logger.info(f"Loaded as {name} with CPU offload")
|
||
return True
|
||
except Exception as e:
|
||
logger.warning(f"{name} + cpu_offload failed: {e}")
|
||
_pipe = None
|
||
_cleanup()
|
||
return False
|
||
|
||
# Try direct device placement
|
||
try:
|
||
_pipe = _pipe.to(target_device)
|
||
logger.info(f"Loaded as {name} on {target_device}")
|
||
return True
|
||
except Exception as e:
|
||
logger.warning(f"{name} + .to({target_device}) failed: {e}")
|
||
_pipe = None
|
||
_cleanup()
|
||
|
||
if not use_offload:
|
||
logger.error(f"{name} could not be placed on {target_device}. On CUDA, try --cpu-offload; on Apple Silicon try a smaller model or lower resolution.")
|
||
return False
|
||
|
||
# OOM — reload and try with CPU offload
|
||
try:
|
||
logger.info(f"Reloading {name} with CPU offload...")
|
||
kwargs = {"torch_dtype": torch_dtype}
|
||
if name == "DiffusionPipeline" and cls_name_from_index and not hasattr(diffusers, cls_name_from_index):
|
||
kwargs["trust_remote_code"] = True
|
||
_pipe = cls.from_pretrained(model_path, **kwargs)
|
||
_fix_meta_tensors(_pipe, torch_dtype)
|
||
_pipe.enable_model_cpu_offload()
|
||
logger.info(f"Loaded as {name} with CPU offload")
|
||
return True
|
||
except Exception as e:
|
||
logger.warning(f"{name} + cpu_offload reload failed: {e}")
|
||
_pipe = None
|
||
_cleanup()
|
||
|
||
# Last resort — sequential offload
|
||
try:
|
||
logger.info(f"Reloading {name} with sequential CPU offload...")
|
||
kwargs = {"torch_dtype": torch_dtype}
|
||
if name == "DiffusionPipeline" and cls_name_from_index and not hasattr(diffusers, cls_name_from_index):
|
||
kwargs["trust_remote_code"] = True
|
||
_pipe = cls.from_pretrained(model_path, **kwargs)
|
||
_fix_meta_tensors(_pipe, torch_dtype)
|
||
_pipe.enable_sequential_cpu_offload()
|
||
logger.info(f"Loaded as {name} with sequential CPU offload")
|
||
return True
|
||
except Exception as e:
|
||
logger.warning(f"{name} + sequential offload failed: {e}")
|
||
_pipe = None
|
||
_cleanup()
|
||
|
||
return False
|
||
|
||
if not loaded:
|
||
for cls, name in candidates:
|
||
if _load_pipe(cls, name):
|
||
loaded = True
|
||
break
|
||
|
||
# Last resort: override unknown pipeline class
|
||
if not loaded and cls_name_from_index and not hasattr(diffusers, cls_name_from_index):
|
||
for fallback in ("Flux2Pipeline", "FluxPipeline", "StableDiffusionPipeline"):
|
||
fb_cls = getattr(diffusers, fallback, None)
|
||
if fb_cls and fb_cls not in [c for c, _ in candidates]:
|
||
logger.info(f"Overriding {cls_name_from_index} -> {fallback}")
|
||
if _load_pipe(fb_cls, fallback):
|
||
loaded = True
|
||
break
|
||
|
||
# Last resort: try from_single_file for raw safetensors / ckpt models
|
||
if not loaded:
|
||
# Find the single-file weight (safetensors preferred, then ckpt/bin)
|
||
single_file = None
|
||
from huggingface_hub import hf_hub_download, list_repo_files
|
||
# Check if it's a HF repo with a single safetensors file
|
||
try:
|
||
files = list_repo_files(model_path)
|
||
sf_files = [f for f in files if f.endswith('.safetensors') and '/' not in f]
|
||
ckpt_files = [f for f in files if f.endswith(('.ckpt', '.bin')) and '/' not in f]
|
||
target = sf_files[0] if sf_files else (ckpt_files[0] if ckpt_files else None)
|
||
if target:
|
||
logger.info(f"Downloading single file: {target}")
|
||
single_file = hf_hub_download(model_path, target)
|
||
except Exception as e:
|
||
logger.warning(f"Could not list repo files for single-file fallback: {e}")
|
||
# Also check local path
|
||
if not single_file:
|
||
local_path = Path(model_path)
|
||
if local_path.is_dir():
|
||
for ext in ('.safetensors', '.ckpt', '.bin'):
|
||
matches = list(local_path.glob(f'*{ext}'))
|
||
if matches:
|
||
single_file = str(matches[0])
|
||
break
|
||
elif local_path.is_file():
|
||
single_file = str(local_path)
|
||
|
||
if single_file:
|
||
logger.info(f"Trying from_single_file with: {single_file}")
|
||
explicit_configs = [
|
||
c.strip()
|
||
for c in str(_args.single_file_config or "").replace("\n", ",").split(",")
|
||
if c.strip()
|
||
]
|
||
config_candidates = explicit_configs or [None]
|
||
_pipeline_configs = [
|
||
("Flux2Pipeline", config_candidates),
|
||
("StableDiffusion3Pipeline", config_candidates),
|
||
("FluxPipeline", config_candidates),
|
||
("StableDiffusionXLPipeline", config_candidates),
|
||
("StableDiffusionPipeline", config_candidates),
|
||
]
|
||
# Deduplicate while preserving order
|
||
_seen = set()
|
||
_deduped = []
|
||
for item in _pipeline_configs:
|
||
if item[0] not in _seen:
|
||
_seen.add(item[0])
|
||
_deduped.append(item)
|
||
_pipeline_configs = _deduped
|
||
# Pre-download config files (json/txt only) so from_single_file doesn't choke
|
||
def _ensure_config_local(repo_id):
|
||
"""Download only config files from a repo, return local path or None."""
|
||
try:
|
||
from huggingface_hub import snapshot_download
|
||
local = snapshot_download(
|
||
repo_id,
|
||
allow_patterns=["*.json", "*.txt", "**/*.json", "**/*.txt"],
|
||
ignore_patterns=["*.safetensors", "*.bin", "*.ckpt", "*.pt", "*.msgpack", "*.h5", "*.onnx", "*.png", "*.jpg", "*.md"],
|
||
token=_hf_token,
|
||
local_files_only=False,
|
||
)
|
||
logger.info(f"Config files cached for {repo_id} at {local}")
|
||
return local
|
||
except Exception as e1:
|
||
logger.warning(f"Could not download configs from {repo_id}: {e1}")
|
||
# Try without allow_patterns (some hf_hub versions have bugs with filters on gated repos)
|
||
try:
|
||
from huggingface_hub import snapshot_download as _sd2
|
||
local = _sd2(
|
||
repo_id,
|
||
ignore_patterns=["*.safetensors", "*.bin", "*.ckpt", "*.pt", "*.msgpack", "*.h5", "*.onnx"],
|
||
token=_hf_token,
|
||
local_files_only=False,
|
||
)
|
||
logger.info(f"Config files cached (no filter) for {repo_id} at {local}")
|
||
return local
|
||
except Exception as e2:
|
||
logger.warning(f"Retry without allow_patterns also failed for {repo_id}: {e2}")
|
||
return None
|
||
|
||
for cls_name, configs in _pipeline_configs:
|
||
if loaded:
|
||
break
|
||
cls = getattr(diffusers, cls_name, None)
|
||
if not cls or not hasattr(cls, 'from_single_file'):
|
||
continue
|
||
for config in configs:
|
||
try:
|
||
kwargs = {"torch_dtype": torch_dtype}
|
||
if config:
|
||
# Use local path instead of repo ID so diffusers doesn't re-download
|
||
local_config = _ensure_config_local(config)
|
||
if not local_config:
|
||
continue
|
||
kwargs["config"] = local_config
|
||
logger.info(f"Trying {cls_name}.from_single_file with config={config}")
|
||
_pipe = cls.from_single_file(single_file, **kwargs)
|
||
_fix_meta_tensors(_pipe, torch_dtype)
|
||
if use_offload and _can_cpu_offload(target_device):
|
||
_pipe.enable_model_cpu_offload()
|
||
logger.info(f"Loaded as {cls_name} (single file, config={config}) with CPU offload")
|
||
else:
|
||
_pipe = _pipe.to(target_device)
|
||
logger.info(f"Loaded as {cls_name} (single file, config={config}) on {target_device}")
|
||
loaded = True
|
||
break
|
||
except Exception as e:
|
||
logger.warning(f"{cls_name}.from_single_file (config={config}) failed: {e}")
|
||
_pipe = None
|
||
_cleanup()
|
||
|
||
if not loaded:
|
||
raise RuntimeError(f"Could not load model from {model_path}. Check diffusers version and model format.")
|
||
|
||
# Memory optimizations
|
||
if _args.attention_slicing:
|
||
try:
|
||
_pipe.enable_attention_slicing()
|
||
logger.info("Attention slicing enabled")
|
||
except Exception:
|
||
pass
|
||
if _args.vae_slicing:
|
||
try:
|
||
_pipe.enable_vae_slicing()
|
||
logger.info("VAE slicing enabled")
|
||
except Exception:
|
||
pass
|
||
|
||
logger.info(f"Model loaded: {_model_id}")
|
||
|
||
# Load LoRA weights if specified
|
||
if _args.lora:
|
||
for lora_path in _args.lora.split(','):
|
||
lora_path = lora_path.strip()
|
||
if not lora_path:
|
||
continue
|
||
try:
|
||
lora_name = Path(lora_path).stem
|
||
_pipe.load_lora_weights(lora_path, adapter_name=lora_name)
|
||
logger.info(f"Loaded LoRA: {lora_name} from {lora_path}")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load LoRA {lora_path}: {e}")
|
||
# Set LoRA scale
|
||
try:
|
||
_pipe.set_adapters([Path(p.strip()).stem for p in _args.lora.split(',') if p.strip()],
|
||
adapter_weights=[_args.lora_scale] * len([p for p in _args.lora.split(',') if p.strip()]))
|
||
logger.info(f"LoRA scale set to {_args.lora_scale}")
|
||
except Exception as e:
|
||
logger.debug(f"Could not set adapter weights: {e}")
|
||
|
||
|
||
@app.get("/v1/models")
|
||
def list_models():
|
||
return {
|
||
"data": [
|
||
{
|
||
"id": _model_id,
|
||
"object": "model",
|
||
"owned_by": "local",
|
||
}
|
||
]
|
||
}
|
||
|
||
|
||
@app.post("/v1/images/generations")
|
||
def generate_image(req: ImageRequest):
|
||
if _pipe is None:
|
||
return {"error": "Model not loaded"}
|
||
|
||
width, height = _parse_size(req.size)
|
||
steps = _quality_steps(req.quality)
|
||
request_id = _start_progress(req.request_id, steps * max(1, int(req.n or 1)), req.prompt, "generation")
|
||
|
||
logger.info(f"Generating: {req.prompt[:80]}... ({width}x{height}, {steps} steps)")
|
||
start = time.time()
|
||
|
||
# Detect if pipeline is inpaint-only (requires image + mask)
|
||
_is_inpaint_pipe = 'inpaint' in type(_pipe).__name__.lower()
|
||
|
||
images = []
|
||
try:
|
||
for image_index in range(req.n):
|
||
progress_offset = image_index * steps
|
||
negative_prompt = _default_negative_prompt() if _pipeline_accepts_arg("negative_prompt") else None
|
||
if _is_inpaint_pipe:
|
||
# Inpaint pipelines need an image + mask — create blank ones for txt2img
|
||
from PIL import Image as _PILGen
|
||
_blank = _PILGen.new('RGB', (width, height), (128, 128, 128))
|
||
_mask = _PILGen.new('L', (width, height), 255) # full white = regenerate everything
|
||
kwargs = {
|
||
"prompt": req.prompt,
|
||
"image": _blank,
|
||
"mask_image": _mask,
|
||
"width": width,
|
||
"height": height,
|
||
"num_inference_steps": steps,
|
||
"guidance_scale": _guidance_scale(),
|
||
}
|
||
else:
|
||
kwargs = {
|
||
"prompt": req.prompt,
|
||
"width": width,
|
||
"height": height,
|
||
"num_inference_steps": steps,
|
||
"guidance_scale": _guidance_scale(),
|
||
}
|
||
if negative_prompt:
|
||
kwargs["negative_prompt"] = negative_prompt
|
||
result = _run_pipeline_with_progress(
|
||
_pipe,
|
||
request_id,
|
||
steps * max(1, int(req.n or 1)),
|
||
**kwargs,
|
||
)
|
||
_update_progress(request_id, progress_offset + steps, steps * max(1, int(req.n or 1)))
|
||
img = result.images[0]
|
||
images.append(img)
|
||
except Exception as e:
|
||
_finish_progress(request_id, "error", str(e))
|
||
raise
|
||
|
||
elapsed = time.time() - start
|
||
logger.info(f"Generated {req.n} image(s) in {elapsed:.1f}s")
|
||
_finish_progress(request_id)
|
||
|
||
return _image_response(images)
|
||
|
||
|
||
@app.get("/v1/images/progress/{request_id}")
|
||
def image_progress(request_id: str):
|
||
item = _PROGRESS.get(request_id)
|
||
if not item:
|
||
return {"id": request_id, "status": "unknown", "step": 0, "total": 0, "percent": 0}
|
||
return item
|
||
|
||
|
||
@app.post("/v1/images/edits")
|
||
async def edit_image(
|
||
prompt: str = Form(...),
|
||
image: UploadFile = File(...),
|
||
model: str = Form(""),
|
||
n: int = Form(1),
|
||
size: str = Form("1024x1024"),
|
||
quality: str = Form("medium"),
|
||
response_format: str = Form("b64_json"),
|
||
request_id: str = Form(""),
|
||
):
|
||
if _pipe is None:
|
||
return {"error": "Model not loaded"}
|
||
accepts_image = _pipeline_accepts_arg("image")
|
||
accepts_input_images = _pipeline_accepts_arg("input_images")
|
||
if not accepts_image and not accepts_input_images:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"{type(_pipe).__name__} does not support image edits. Use /v1/images/generations with this model.",
|
||
)
|
||
|
||
from PIL import Image as PILImage, ImageOps
|
||
|
||
width, height = _parse_size(size)
|
||
steps = _quality_steps(quality)
|
||
request_id = _start_progress(request_id, steps * max(1, min(int(n or 1), 4)), prompt, "edit")
|
||
raw = await image.read()
|
||
init_image = PILImage.open(io.BytesIO(raw)).convert("RGB")
|
||
if width > 0 and height > 0:
|
||
init_image = ImageOps.fit(init_image, (width, height), method=PILImage.LANCZOS, centering=(0.5, 0.5))
|
||
|
||
logger.info(f"Editing image: {prompt[:80]}... ({width}x{height}, {steps} steps)")
|
||
start = time.time()
|
||
images = []
|
||
total_images = max(1, min(int(n or 1), 4))
|
||
for image_index in range(total_images):
|
||
progress_offset = image_index * steps
|
||
try:
|
||
if accepts_input_images and not accepts_image:
|
||
negative_prompt = _default_negative_prompt()
|
||
kwargs = _pipeline_call_kwargs(
|
||
prompt=prompt,
|
||
input_images=[init_image],
|
||
width=width,
|
||
height=height,
|
||
num_inference_steps=steps,
|
||
max_sequence_length=1024,
|
||
text_guidance_scale=_guidance_scale(),
|
||
image_guidance_scale=2.0,
|
||
cfg_range=(0.0, 1.0),
|
||
negative_prompt=negative_prompt,
|
||
num_images_per_prompt=1,
|
||
output_type="pil",
|
||
max_pixels=width * height if width > 0 and height > 0 else None,
|
||
max_input_image_side_length=max(width, height) if width > 0 and height > 0 else None,
|
||
)
|
||
else:
|
||
kwargs = _pipeline_call_kwargs(
|
||
image=init_image,
|
||
prompt=prompt,
|
||
width=width,
|
||
height=height,
|
||
num_inference_steps=steps,
|
||
guidance_scale=3.5,
|
||
true_cfg_scale=4.0,
|
||
negative_prompt=_default_negative_prompt(),
|
||
output_type="pil",
|
||
)
|
||
result = _run_pipeline_with_progress(
|
||
_pipe,
|
||
request_id,
|
||
steps * total_images,
|
||
**kwargs,
|
||
)
|
||
except TypeError:
|
||
if accepts_input_images and not accepts_image:
|
||
kwargs = _pipeline_call_kwargs(
|
||
prompt=prompt,
|
||
input_images=[init_image],
|
||
num_inference_steps=steps,
|
||
text_guidance_scale=_guidance_scale(),
|
||
image_guidance_scale=2.0,
|
||
negative_prompt=_default_negative_prompt(),
|
||
output_type="pil",
|
||
)
|
||
else:
|
||
kwargs = _pipeline_call_kwargs(
|
||
image=init_image,
|
||
prompt=prompt,
|
||
num_inference_steps=steps,
|
||
guidance_scale=3.5,
|
||
negative_prompt=_default_negative_prompt(),
|
||
output_type="pil",
|
||
)
|
||
result = _run_pipeline_with_progress(
|
||
_pipe,
|
||
request_id,
|
||
steps * total_images,
|
||
**kwargs,
|
||
)
|
||
except Exception as e:
|
||
_finish_progress(request_id, "error", str(e))
|
||
raise
|
||
_update_progress(request_id, progress_offset + steps, steps * total_images)
|
||
images.append(result.images[0])
|
||
|
||
elapsed = time.time() - start
|
||
logger.info(f"Edited {len(images)} image(s) in {elapsed:.1f}s")
|
||
_finish_progress(request_id)
|
||
return _image_response(images)
|
||
|
||
|
||
class InpaintRequest(BaseModel):
|
||
image: str # base64 PNG
|
||
mask: str # base64 PNG (white = inpaint area)
|
||
prompt: str
|
||
width: int = 0
|
||
height: int = 0
|
||
steps: int = 0
|
||
strength: float = 0.75 # how much to change (0=nothing, 1=full regeneration)
|
||
feather: int = 8 # mask edge feathering in pixels
|
||
|
||
|
||
_inpaint_pipe = None
|
||
_img2img_pipe = None
|
||
|
||
def _get_inpaint_pipe():
|
||
"""Lazy-load an inpaint or img2img pipeline from the same model."""
|
||
global _inpaint_pipe, _img2img_pipe
|
||
if _inpaint_pipe:
|
||
return _inpaint_pipe, 'inpaint'
|
||
if _img2img_pipe:
|
||
return _img2img_pipe, 'img2img'
|
||
|
||
import diffusers
|
||
model_path = _args.model
|
||
torch_dtype = DTYPE_MAP.get(_args.dtype, torch.bfloat16)
|
||
|
||
# Check if the main pipeline IS already an inpaint pipeline
|
||
pipe_cls_name = type(_pipe).__name__
|
||
if 'inpaint' in pipe_cls_name.lower():
|
||
_inpaint_pipe = _pipe
|
||
logger.info(f"Main pipeline is already inpaint: {pipe_cls_name}")
|
||
# Also try to get img2img from it
|
||
try:
|
||
img2img_cls_name = pipe_cls_name.replace('Inpaint', 'Img2Img')
|
||
img2img_cls = getattr(diffusers, img2img_cls_name, None)
|
||
if img2img_cls:
|
||
_img2img_pipe = img2img_cls.from_pipe(_pipe)
|
||
logger.info(f"Also loaded img2img from inpaint pipe: {img2img_cls_name}")
|
||
except Exception as e:
|
||
logger.debug(f"Could not create img2img from inpaint: {e}")
|
||
return _inpaint_pipe, 'inpaint'
|
||
|
||
# Try loading a dedicated inpaint pipeline from the same components
|
||
inpaint_names = [
|
||
pipe_cls_name.replace('Pipeline', 'InpaintPipeline'),
|
||
'StableDiffusion3InpaintPipeline',
|
||
'StableDiffusionXLInpaintPipeline',
|
||
'StableDiffusionInpaintPipeline',
|
||
]
|
||
for name in inpaint_names:
|
||
cls = getattr(diffusers, name, None)
|
||
if cls:
|
||
try:
|
||
_inpaint_pipe = cls.from_pipe(_pipe)
|
||
logger.info(f"Loaded inpaint pipeline: {name}")
|
||
return _inpaint_pipe, 'inpaint'
|
||
except Exception as e:
|
||
logger.debug(f"{name} from_pipe failed: {e}")
|
||
|
||
# Try img2img pipeline
|
||
img2img_names = [
|
||
pipe_cls_name.replace('Pipeline', 'Img2ImgPipeline'),
|
||
'StableDiffusion3Img2ImgPipeline',
|
||
'StableDiffusionXLImg2ImgPipeline',
|
||
'StableDiffusionImg2ImgPipeline',
|
||
]
|
||
torch_dtype = DTYPE_MAP.get(_args.dtype, torch.bfloat16)
|
||
harmonize_gpu = _args.harmonize_gpu
|
||
target_device = _target_device()
|
||
for name in img2img_names:
|
||
cls = getattr(diffusers, name, None)
|
||
if cls:
|
||
try:
|
||
if harmonize_gpu is not None and target_device == "cuda":
|
||
# Load fresh on separate GPU
|
||
logger.info(f"Loading {name} on cuda:{harmonize_gpu}...")
|
||
_img2img_pipe = cls.from_pretrained(_args.model, torch_dtype=torch_dtype)
|
||
_img2img_pipe = _img2img_pipe.to(f"cuda:{harmonize_gpu}")
|
||
else:
|
||
_img2img_pipe = cls.from_pipe(_pipe, torch_dtype=torch_dtype)
|
||
logger.info(f"Loaded img2img pipeline: {name}" + (f" on cuda:{harmonize_gpu}" if harmonize_gpu is not None else ""))
|
||
return _img2img_pipe, 'img2img'
|
||
except Exception as e:
|
||
logger.debug(f"{name} failed: {e}")
|
||
try:
|
||
# Some pipelines need from_pretrained instead of from_pipe
|
||
_img2img_pipe = cls.from_pretrained(_args.model, torch_dtype=torch_dtype)
|
||
if _args.cpu_offload and _can_cpu_offload(target_device):
|
||
_img2img_pipe.enable_model_cpu_offload()
|
||
else:
|
||
_img2img_pipe = _img2img_pipe.to(target_device)
|
||
logger.info(f"Loaded img2img pipeline (from_pretrained): {name}")
|
||
return _img2img_pipe, 'img2img'
|
||
except Exception as e2:
|
||
logger.debug(f"{name} from_pretrained also failed: {e2}")
|
||
|
||
logger.warning("No inpaint or img2img pipeline available — will use txt2img fallback")
|
||
return None, None
|
||
|
||
|
||
@app.post("/v1/images/inpaint")
|
||
def inpaint_image(req: InpaintRequest):
|
||
"""Inpaint masked region. Tries: native inpaint → img2img+composite → txt2img+composite."""
|
||
if _pipe is None:
|
||
return {"error": "Model not loaded"}
|
||
|
||
from PIL import Image as PILImage
|
||
|
||
# Decode input image and mask
|
||
img_bytes = base64.b64decode(req.image)
|
||
mask_bytes = base64.b64decode(req.mask)
|
||
init_image = PILImage.open(io.BytesIO(img_bytes)).convert("RGB")
|
||
mask_image = PILImage.open(io.BytesIO(mask_bytes)).convert("L")
|
||
|
||
# Feather value — applied after cropping to avoid edge clipping
|
||
feather = max(0, min(60, req.feather))
|
||
|
||
width = req.width or init_image.width
|
||
height = req.height or init_image.height
|
||
|
||
default_steps = _args.steps or 12
|
||
steps = req.steps or default_steps
|
||
|
||
logger.info(f"Inpainting: {req.prompt[:80]}... ({width}x{height}, {steps} steps)")
|
||
start = time.time()
|
||
|
||
strength = max(0.1, min(1.0, req.strength))
|
||
|
||
# Try to get a dedicated inpaint or img2img pipeline
|
||
alt_pipe, alt_type = _get_inpaint_pipe()
|
||
|
||
# SDXL inpaint expects ~1024 on the short side. Running at canvas
|
||
# native resolution can produce grey / muted output when the model's
|
||
# latent grid is far larger than what it was trained on. Cap to a
|
||
# model-friendly box (multiples of 8), inpaint there, upscale back.
|
||
max_side = 1024
|
||
scale = min(max_side / max(width, height), 1.0)
|
||
work_w = max(64, ((int(width * scale) + 7) // 8) * 8)
|
||
work_h = max(64, ((int(height * scale) + 7) // 8) * 8)
|
||
work_init = init_image.resize((work_w, work_h), PILImage.LANCZOS)
|
||
work_mask = mask_image.resize((work_w, work_h), PILImage.BILINEAR)
|
||
logger.info(f"Inpaint working size: {work_w}x{work_h} (from {width}x{height})")
|
||
|
||
# SDXL VAE in fp16/bfloat16 commonly produces NaN/overflow that
|
||
# decodes to flat grey output. Upcast the VAE to fp32 before the
|
||
# call; cheap (only the VAE decode runs in fp32, the heavy UNet
|
||
# stays in the requested dtype). One-time per pipeline.
|
||
if alt_pipe is not None and not getattr(alt_pipe, '_ge_vae_upcast', False):
|
||
try:
|
||
alt_pipe.upcast_vae()
|
||
alt_pipe._ge_vae_upcast = True
|
||
logger.info("Upcast VAE to fp32 to avoid grey-output bug")
|
||
except Exception as e:
|
||
logger.warning(f"Could not upcast VAE: {e}")
|
||
|
||
try:
|
||
if alt_type == 'inpaint' and alt_pipe:
|
||
# Use dedicated inpaint pipeline. guidance_scale 7.5 is the
|
||
# SDXL default — the previous 3.5 was producing muted / grey
|
||
# results, especially on style-transfer prompts with large
|
||
# masks.
|
||
logger.info("Using dedicated inpaint pipeline")
|
||
result = alt_pipe(
|
||
prompt=req.prompt,
|
||
image=work_init,
|
||
mask_image=work_mask,
|
||
width=work_w,
|
||
height=work_h,
|
||
num_inference_steps=steps,
|
||
strength=strength,
|
||
guidance_scale=7.5,
|
||
)
|
||
elif alt_type == 'img2img' and alt_pipe:
|
||
raise TypeError("Skip to img2img fallback")
|
||
else:
|
||
# Try the main pipeline with inpaint args
|
||
result = _pipe(
|
||
prompt=req.prompt,
|
||
image=work_init,
|
||
mask_image=work_mask,
|
||
width=work_w,
|
||
height=work_h,
|
||
num_inference_steps=steps,
|
||
strength=strength,
|
||
guidance_scale=7.5,
|
||
)
|
||
except TypeError:
|
||
# Pipeline doesn't support native inpainting — use crop-to-mask + img2img + composite
|
||
# This preserves context by only regenerating the masked region with surrounding padding
|
||
import numpy as np
|
||
logger.info(f"Pipeline doesn't support inpainting — using crop+img2img (strength={strength}) + composite")
|
||
|
||
mask_resized = mask_image.resize((width, height))
|
||
init_resized = init_image.resize((width, height))
|
||
mask_arr = np.array(mask_resized)
|
||
|
||
# Find bounding box of the mask
|
||
ys, xs = np.where(mask_arr > 10)
|
||
if len(xs) == 0 or len(ys) == 0:
|
||
logger.warning("Empty mask — returning original image")
|
||
buf = io.BytesIO()
|
||
init_resized.save(buf, format="PNG")
|
||
return {"image": base64.b64encode(buf.getvalue()).decode(), "elapsed": 0}
|
||
|
||
x1, y1, x2, y2 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
|
||
|
||
# Add generous padding (50% of mask size, min 64px) so model sees surrounding context
|
||
pad_x = max(64, int((x2 - x1) * 0.5))
|
||
pad_y = max(64, int((y2 - y1) * 0.5))
|
||
cx1 = max(0, x1 - pad_x)
|
||
cy1 = max(0, y1 - pad_y)
|
||
cx2 = min(width, x2 + pad_x)
|
||
cy2 = min(height, y2 + pad_y)
|
||
|
||
# Make crop square and round to multiple of 64 (SD3 VAE requirement)
|
||
crop_size = max(cx2 - cx1, cy2 - cy1)
|
||
crop_size = max(256, ((crop_size + 63) // 64) * 64) # min 256, round up to 64
|
||
# Center the square crop on the mask center
|
||
cx_mid = (cx1 + cx2) // 2
|
||
cy_mid = (cy1 + cy2) // 2
|
||
cx1 = max(0, cx_mid - crop_size // 2)
|
||
cy1 = max(0, cy_mid - crop_size // 2)
|
||
cx2 = min(width, cx1 + crop_size)
|
||
cy2 = min(height, cy1 + crop_size)
|
||
# Adjust if we hit image edges
|
||
if cx2 - cx1 < crop_size:
|
||
cx1 = max(0, cx2 - crop_size)
|
||
if cy2 - cy1 < crop_size:
|
||
cy1 = max(0, cy2 - crop_size)
|
||
cw = cx2 - cx1
|
||
ch = cy2 - cy1
|
||
|
||
logger.info(f"Mask bbox: ({x1},{y1})-({x2},{y2}), crop region: ({cx1},{cy1})-({cx2},{cy2}) = {cw}x{ch}")
|
||
|
||
# Crop the original image and mask to the region
|
||
crop_img = init_resized.crop((cx1, cy1, cx2, cy2))
|
||
crop_mask = mask_resized.crop((cx1, cy1, cx2, cy2))
|
||
|
||
# Use img2img pipeline if available, otherwise fall back
|
||
_i2i_pipe = alt_pipe if alt_type == 'img2img' else None
|
||
# Ensure crop image is properly sized (multiple of 8)
|
||
crop_img = crop_img.resize((cw, ch))
|
||
try:
|
||
if _i2i_pipe:
|
||
logger.info(f"Using img2img pipeline on crop ({cw}x{ch})")
|
||
result = _i2i_pipe(
|
||
prompt=req.prompt,
|
||
image=crop_img,
|
||
num_inference_steps=steps,
|
||
strength=strength,
|
||
guidance_scale=7.0,
|
||
)
|
||
else:
|
||
# Try main pipeline with image arg
|
||
result = _pipe(
|
||
prompt=req.prompt,
|
||
image=crop_img,
|
||
num_inference_steps=steps,
|
||
strength=strength,
|
||
guidance_scale=3.5,
|
||
)
|
||
generated_crop = result.images[0].resize((cw, ch))
|
||
except TypeError:
|
||
# No img2img support at all — txt2img on crop size
|
||
logger.info("No img2img support — txt2img on crop region")
|
||
result = _pipe(
|
||
prompt=req.prompt,
|
||
width=cw,
|
||
height=ch,
|
||
num_inference_steps=steps,
|
||
guidance_scale=3.5,
|
||
)
|
||
generated_crop = result.images[0].resize((cw, ch))
|
||
|
||
# Apply feathering to the cropped mask for soft blending edges
|
||
if feather > 0:
|
||
from PIL import ImageFilter
|
||
# PIL GaussianBlur radius is ~half of CSS blur pixels, so multiply
|
||
blur_radius = feather * 1.5
|
||
crop_mask = crop_mask.filter(ImageFilter.GaussianBlur(radius=blur_radius))
|
||
logger.info(f"Applied {feather}px feather (PIL radius={blur_radius:.0f}) to crop mask")
|
||
|
||
# Composite: blend generated crop into original using the feathered mask
|
||
orig_arr = np.array(init_resized).astype(float)
|
||
gen_full = orig_arr.copy()
|
||
crop_gen_arr = np.array(generated_crop).astype(float)
|
||
crop_mask_arr = np.array(crop_mask) / 255.0
|
||
|
||
# Blend only in the crop region
|
||
region = gen_full[cy1:cy2, cx1:cx2]
|
||
blended_region = region * (1 - crop_mask_arr[:, :, None]) + crop_gen_arr * crop_mask_arr[:, :, None]
|
||
gen_full[cy1:cy2, cx1:cx2] = blended_region
|
||
|
||
result_img = PILImage.fromarray(gen_full.astype(np.uint8))
|
||
|
||
buf = io.BytesIO()
|
||
result_img.save(buf, format="PNG")
|
||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||
elapsed = time.time() - start
|
||
logger.info(f"Inpaint (crop+composite) done in {elapsed:.1f}s")
|
||
return {"image": b64, "elapsed": round(elapsed, 2)}
|
||
|
||
img = result.images[0]
|
||
# Upscale back to the canvas size if we worked at a smaller resolution.
|
||
if (img.width, img.height) != (width, height):
|
||
img = img.resize((width, height), PILImage.LANCZOS)
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||
elapsed = time.time() - start
|
||
logger.info(f"Inpaint done in {elapsed:.1f}s")
|
||
return {"image": b64, "elapsed": round(elapsed, 2)}
|
||
|
||
|
||
class HarmonizeRequest(BaseModel):
|
||
image: str # base64 PNG
|
||
prompt: str
|
||
# Two-stage harmonize:
|
||
# 1) Reinhard color transfer inside `body_mask` (matches L*a*b* mean/std
|
||
# of the masked region to the unmasked surroundings). Pixel-sharp.
|
||
# 2) Optional narrow inpaint on `seam_mask` (alpha edge band) to fix
|
||
# jagged cutouts and seams. Only the edge band is regenerated.
|
||
color_match: float = 0.65 # 0..1 — how much of the color shift to apply
|
||
seam_fix: float = 0.0 # 0..1 — strength of the seam inpaint pass
|
||
body_mask: str | None = None # base64 PNG, white = layer body
|
||
seam_mask: str | None = None # base64 PNG, white = layer alpha edge band
|
||
steps: int = 0
|
||
# Legacy fields (older clients): if `mask` is sent without body/seam,
|
||
# we treat it as body_mask. `strength` maps to color_match.
|
||
mask: str | None = None
|
||
strength: float | None = None
|
||
max_side: int = 1024
|
||
|
||
|
||
def _rgb_to_lalphabeta(rgb_f):
|
||
"""RGB → L*alpha*beta (Ruderman et al., the colour space Reinhard's
|
||
original paper used). Pure numpy — no cv2. Input/output: float32 arrays
|
||
of shape (..., 3); input in 0..255, output unbounded log-RGB-style."""
|
||
import numpy as np
|
||
eps = 1.0
|
||
# Linearise to LMS cone space
|
||
M_rgb2lms = np.array([
|
||
[0.3811, 0.5783, 0.0402],
|
||
[0.1967, 0.7244, 0.0782],
|
||
[0.0241, 0.1288, 0.8444],
|
||
], dtype=np.float32)
|
||
lms = rgb_f @ M_rgb2lms.T
|
||
lms = np.log(np.maximum(lms, eps))
|
||
# LMS → L*alpha*beta
|
||
M_lms2lab = np.array([
|
||
[1.0/np.sqrt(3), 1.0/np.sqrt(3), 1.0/np.sqrt(3)],
|
||
[1.0/np.sqrt(6), 1.0/np.sqrt(6), -2.0/np.sqrt(6)],
|
||
[1.0/np.sqrt(2), -1.0/np.sqrt(2), 0.0 ],
|
||
], dtype=np.float32)
|
||
return lms @ M_lms2lab.T
|
||
|
||
|
||
def _lalphabeta_to_rgb(lab):
|
||
"""Inverse of _rgb_to_lalphabeta. Returns RGB float32 in 0..255 (clipped)."""
|
||
import numpy as np
|
||
M_lab2lms = np.array([
|
||
[np.sqrt(3)/3.0, np.sqrt(6)/6.0, np.sqrt(2)/2.0],
|
||
[np.sqrt(3)/3.0, np.sqrt(6)/6.0, -np.sqrt(2)/2.0],
|
||
[np.sqrt(3)/3.0, -np.sqrt(6)/3.0, 0.0 ],
|
||
], dtype=np.float32)
|
||
lms = lab @ M_lab2lms.T
|
||
lms = np.exp(lms)
|
||
M_lms2rgb = np.array([
|
||
[ 4.4679, -3.5873, 0.1193],
|
||
[-1.2186, 2.3809, -0.1624],
|
||
[ 0.0497, -0.2439, 1.2045],
|
||
], dtype=np.float32)
|
||
rgb = lms @ M_lms2rgb.T
|
||
return np.clip(rgb, 0, 255)
|
||
|
||
|
||
def _reinhard_color_transfer(source_rgb, body_mask_l, blend: float = 1.0):
|
||
"""Match the masked region's color statistics to the unmasked
|
||
surroundings using Reinhard's L*alpha*beta transfer. Pure numpy.
|
||
|
||
`blend` (0..1) controls how much of the shift to apply.
|
||
"""
|
||
import numpy as np
|
||
from PIL import Image as _PILImg
|
||
|
||
src_np = np.asarray(source_rgb).astype(np.float32) # H,W,3 in 0..255
|
||
h, w, _ = src_np.shape
|
||
|
||
mask_np = np.asarray(body_mask_l).astype(np.float32) / 255.0
|
||
if mask_np.shape != (h, w):
|
||
return source_rgb
|
||
|
||
interior = mask_np > 0.5
|
||
exterior = mask_np < 0.05
|
||
if interior.sum() < 100 or exterior.sum() < 100:
|
||
return source_rgb
|
||
|
||
lab = _rgb_to_lalphabeta(src_np)
|
||
in_pix = lab[interior]
|
||
out_pix = lab[exterior]
|
||
|
||
in_mean, in_std = in_pix.mean(axis=0), in_pix.std(axis=0) + 1e-6
|
||
out_mean, out_std = out_pix.mean(axis=0), out_pix.std(axis=0) + 1e-6
|
||
|
||
shifted = lab.copy()
|
||
shifted[interior] = (lab[interior] - in_mean) * (out_std / in_std) + out_mean
|
||
rgb_shifted = _lalphabeta_to_rgb(shifted)
|
||
|
||
# Lerp source ↔ shifted, weighted by mask × blend so the edge of the
|
||
# mask fades back to source smoothly.
|
||
m3 = (mask_np * blend)[..., None]
|
||
out = src_np * (1 - m3) + rgb_shifted * m3
|
||
return _PILImg.fromarray(np.clip(out, 0, 255).astype(np.uint8), mode='RGB')
|
||
|
||
|
||
def _decode_mask_b64(b64_str, target_size):
|
||
"""Decode a base64-encoded grayscale PNG. Returns PIL 'L' image at
|
||
`target_size`, or None if empty/invalid."""
|
||
if not b64_str:
|
||
return None
|
||
try:
|
||
from PIL import Image as _PILImg
|
||
m = _PILImg.open(io.BytesIO(base64.b64decode(b64_str))).convert("L")
|
||
if m.size != target_size:
|
||
m = m.resize(target_size, _PILImg.BILINEAR)
|
||
if not m.getbbox():
|
||
return None
|
||
return m
|
||
except Exception as e:
|
||
logger.warning(f"Harmonize: bad mask: {e}")
|
||
return None
|
||
|
||
|
||
@app.post("/v1/images/harmonize")
|
||
def harmonize_image(req: HarmonizeRequest):
|
||
"""Two-stage layer harmonization.
|
||
|
||
Stage 1 — Reinhard color transfer inside `body_mask`: matches the
|
||
masked region's L*a*b* mean/std to the unmasked surroundings. Pixel-
|
||
sharp, no model regen. Controlled by `color_match` (0..1).
|
||
|
||
Stage 2 — Optional narrow inpaint on `seam_mask` (alpha edge band):
|
||
only the band is regenerated; layer interiors stay identical to the
|
||
color-shifted result. Controlled by `seam_fix` (0..1). Skipped if
|
||
`seam_fix=0` or no inpaint pipeline is available.
|
||
|
||
Backwards compat: if only `mask` is provided (no body/seam), it's
|
||
treated as body_mask. `strength` (old field) maps to `color_match`.
|
||
"""
|
||
if _pipe is None:
|
||
return {"error": "Model not loaded"}
|
||
|
||
from PIL import Image as PILImage
|
||
|
||
img_bytes = base64.b64decode(req.image)
|
||
source_full = PILImage.open(io.BytesIO(img_bytes)).convert("RGB")
|
||
orig_w, orig_h = source_full.size
|
||
|
||
# Resolve old-vs-new field names.
|
||
body_b64 = req.body_mask or req.mask
|
||
seam_b64 = req.seam_mask
|
||
color_match = req.color_match
|
||
if req.strength is not None:
|
||
color_match = req.strength
|
||
color_match = max(0.0, min(1.0, color_match))
|
||
seam_fix = max(0.0, min(1.0, req.seam_fix))
|
||
|
||
body_mask_full = _decode_mask_b64(body_b64, (orig_w, orig_h))
|
||
seam_mask_full = _decode_mask_b64(seam_b64, (orig_w, orig_h))
|
||
|
||
# If neither mask was supplied: legacy whole-image fallback. The user
|
||
# didn't tell us where the seams are, so we can't do targeted blending.
|
||
if body_mask_full is None and seam_mask_full is None:
|
||
logger.info("Harmonize: no masks — falling back to legacy whole-image path")
|
||
return _legacy_whole_image_harmonize(req, source_full)
|
||
|
||
logger.info(
|
||
f"Harmonize: color_match={color_match:.2f} seam_fix={seam_fix:.2f} "
|
||
f"body_mask={'y' if body_mask_full else 'n'} seam_mask={'y' if seam_mask_full else 'n'}"
|
||
)
|
||
start = time.time()
|
||
|
||
# ── Stage 1: Reinhard color transfer (pixel-sharp, no regen) ──
|
||
if body_mask_full is not None and color_match > 0.01:
|
||
try:
|
||
stage1 = _reinhard_color_transfer(source_full, body_mask_full, blend=color_match)
|
||
except Exception as e:
|
||
logger.warning(f"Harmonize stage 1 failed, skipping: {e}")
|
||
stage1 = source_full
|
||
else:
|
||
stage1 = source_full
|
||
|
||
# ── Stage 2: narrow seam inpaint (only the alpha edge band) ──
|
||
final = stage1
|
||
if seam_mask_full is not None and seam_fix > 0.01:
|
||
alt_pipe, alt_type = _get_inpaint_pipe()
|
||
is_inpaint_main = 'inpaint' in type(_pipe).__name__.lower()
|
||
inpaint_pipe = alt_pipe if alt_type == 'inpaint' else (_pipe if is_inpaint_main else None)
|
||
if inpaint_pipe is None:
|
||
logger.info("Harmonize: seam_fix requested but no inpaint pipe — skipping stage 2")
|
||
else:
|
||
try:
|
||
max_side = req.max_side or 1024
|
||
scale = min(max_side / orig_w, max_side / orig_h, 1.0)
|
||
w = ((int(orig_w * scale) + 63) // 64) * 64
|
||
h = ((int(orig_h * scale) + 63) // 64) * 64
|
||
init_small = stage1.resize((w, h), PILImage.LANCZOS)
|
||
seam_small = seam_mask_full.resize((w, h), PILImage.BILINEAR)
|
||
# Cap the inpaint strength — seam_fix=1.0 → strength=0.50,
|
||
# so even max setting can't fully redraw the band.
|
||
inpaint_strength = max(0.10, min(0.50, seam_fix * 0.50))
|
||
steps = req.steps or (_args.steps or 12)
|
||
logger.info(f"Harmonize stage 2: seam inpaint at {w}x{h}, strength={inpaint_strength:.2f}")
|
||
result = inpaint_pipe(
|
||
prompt=req.prompt,
|
||
image=init_small,
|
||
mask_image=seam_small,
|
||
width=w,
|
||
height=h,
|
||
num_inference_steps=max(steps, 20),
|
||
strength=inpaint_strength,
|
||
guidance_scale=7.0,
|
||
)
|
||
ai_small = result.images[0]
|
||
ai_full = ai_small.resize((orig_w, orig_h), PILImage.LANCZOS) if (w, h) != (orig_w, orig_h) else ai_small
|
||
# Composite back using the seam mask as alpha — outside the
|
||
# seam band stays pixel-identical to stage1.
|
||
final = PILImage.composite(ai_full, stage1, seam_mask_full)
|
||
except Exception as e:
|
||
logger.warning(f"Harmonize stage 2 failed, returning stage 1 only: {e}")
|
||
final = stage1
|
||
|
||
buf = io.BytesIO()
|
||
final.save(buf, format="PNG")
|
||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||
elapsed = time.time() - start
|
||
logger.info(f"Harmonize done in {elapsed:.1f}s")
|
||
return {"image": b64, "elapsed": round(elapsed, 2)}
|
||
|
||
|
||
def _legacy_whole_image_harmonize(req, source_full):
|
||
"""Old behaviour: no masks supplied → run img2img on the entire image.
|
||
Kept for cases where the client wants a global re-render."""
|
||
from PIL import Image as PILImage
|
||
|
||
orig_w, orig_h = source_full.size
|
||
max_side = req.max_side or 1024
|
||
scale = min(max_side / orig_w, max_side / orig_h, 1.0)
|
||
width = ((int(orig_w * scale) + 63) // 64) * 64
|
||
height = ((int(orig_h * scale) + 63) // 64) * 64
|
||
init_image = source_full.resize((width, height), PILImage.LANCZOS)
|
||
steps = req.steps or (_args.steps or 12)
|
||
strength = req.strength if req.strength is not None else 0.30
|
||
strength = max(0.1, min(0.9, strength))
|
||
|
||
alt_pipe, alt_type = _get_inpaint_pipe()
|
||
i2i_pipe = _img2img_pipe if _img2img_pipe else (alt_pipe if alt_type == 'img2img' else None)
|
||
|
||
start = time.time()
|
||
try:
|
||
if i2i_pipe:
|
||
result = i2i_pipe(
|
||
prompt=req.prompt, image=init_image,
|
||
num_inference_steps=steps, strength=strength, guidance_scale=7.0,
|
||
)
|
||
else:
|
||
result = _pipe(
|
||
prompt=req.prompt, image=init_image,
|
||
num_inference_steps=steps, strength=strength, guidance_scale=7.0,
|
||
)
|
||
except TypeError:
|
||
result = _pipe(
|
||
prompt=req.prompt, width=width, height=height,
|
||
num_inference_steps=steps, guidance_scale=7.0,
|
||
)
|
||
|
||
img = result.images[0]
|
||
if (orig_w, orig_h) != (width, height):
|
||
img = img.resize((orig_w, orig_h), PILImage.LANCZOS)
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||
elapsed = time.time() - start
|
||
logger.info(f"Legacy harmonize done in {elapsed:.1f}s")
|
||
return {"image": b64, "elapsed": round(elapsed, 2)}
|
||
|
||
|
||
@app.get("/health")
|
||
def health():
|
||
return {"status": "ok", "model": _model_id}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--model", required=True, help="Path to diffusers model")
|
||
parser.add_argument("--lora", type=str, default=None, help="Path to LoRA weights (.safetensors). Can specify multiple comma-separated.")
|
||
parser.add_argument("--lora-scale", type=float, default=1.0, help="LoRA weight scale (0.0-2.0)")
|
||
parser.add_argument("--port", type=int, default=8100)
|
||
parser.add_argument("--host", default="127.0.0.1")
|
||
parser.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
|
||
parser.add_argument("--device-map", default=None, help="Device map strategy (unused, kept for compat)")
|
||
parser.add_argument("--steps", type=int, default=0, help="Default inference steps (0=auto)")
|
||
parser.add_argument("--guidance-scale", type=float, default=3.5, help="Default classifier-free guidance scale")
|
||
parser.add_argument("--negative-prompt", default="", help="Default negative prompt for pipelines that support it")
|
||
parser.add_argument("--single-file-config", default="", help="Base Diffusers repo/path for single-file checkpoints that need missing components. Comma-separated values are tried in order.")
|
||
parser.add_argument("--width", type=int, default=1024, help="Default output width")
|
||
parser.add_argument("--height", type=int, default=1024, help="Default output height")
|
||
parser.add_argument("--cpu-offload", action="store_true", help="Enable model CPU offload")
|
||
parser.add_argument("--attention-slicing", action="store_true", help="Enable attention slicing")
|
||
parser.add_argument("--vae-slicing", action="store_true", help="Enable VAE slicing")
|
||
parser.add_argument("--harmonize-gpu", type=int, default=None, help="GPU index for harmonize/img2img (default: same as main)")
|
||
parser.add_argument("--allowed-host", action="append", default=[],
|
||
help="Additional Host header value to accept (DNS-rebinding allowlist). "
|
||
"Can be repeated. Loopback values are always included.")
|
||
parser.add_argument("--allowed-origin", action="append", default=[],
|
||
help="Additional CORS origin to allow. Can be repeated. Defaults to "
|
||
"no cross-origin access — only pass this if you need a browser "
|
||
"on a specific origin to call the server.")
|
||
_args = parser.parse_args()
|
||
|
||
# Replace the module-load middleware stack with the CLI-configured one so
|
||
# operator-supplied --allowed-host / --allowed-origin values take effect
|
||
# before the first request is served. user_middleware is consulted lazily
|
||
# when the middleware stack is built on the first request, so mutating it
|
||
# here is safe.
|
||
final_hosts = _compute_allowed_hosts(_args.host, _args.allowed_host)
|
||
final_origins = _compute_cors_origins(_args.allowed_origin)
|
||
_configure_security_middleware(app, final_hosts, final_origins)
|
||
logger.info("security middleware: allowed_hosts=%s allowed_origins=%s",
|
||
final_hosts, final_origins or "(none — default-deny)")
|
||
|
||
app.state.model_path = _args.model
|
||
uvicorn.run(app, host=_args.host, port=_args.port)
|