#!/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)