#!/usr/bin/env python3 """ OpenAI-compatible API wrapper for ComfyUI. Supports: ideogram4, flux2 Endpoint: POST /v1/images/generations Extended params: seed, steps, cfg, quality, negative_prompt, style Timeout: 1800s (30 min) """ import os import re import json import base64 import time import asyncio import httpx from typing import Optional from fastapi import FastAPI, Request, HTTPException, Header from fastapi.responses import JSONResponse from pydantic import BaseModel, Field COMFY_HOST = os.environ.get("COMFYUI_HOST", "localhost") COMFY_PORT = os.environ.get("COMFYUI_PORT", "8188") COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}" API_KEY = os.environ.get("API_KEY", "") ALLOWED_SIZES = { "256x256", "512x512", "512x768", "768x512", "768x1024", "1024x768", "1024x1024", "1024x1536", "1536x1024", "832x1216", "1216x832", "896x1152", "1152x896", "640x1536", "1536x640" } DEFAULT_SIZE = "1024x1024" QUALITY_STEPS = { "low": 8, "medium": 20, "high": 30, "ultra": 50, } IDEOGRAM_ALLOWED = {"1024x1024", "1024x1536", "1536x1024", "832x1216", "1216x832", "896x1152", "1152x896"} IDEOGRAM_DEFAULT = "1024x1024" FLUX_ALLOWED = {"256x256", "512x512", "512x768", "768x512", "768x1024", "1024x768", "1024x1024", "1024x1536", "1536x1024", "832x1216", "1216x832"} # Only two models active on this system ALLOWED_MODELS = {"ideogram4", "flux2"} DEFAULT_MODEL = "ideogram4" def _parse_size(size: str) -> tuple[int, int]: m = re.match(r"(\d+)x(\d+)", size) if not m: raise ValueError(f"Invalid size: {size}") return int(m.group(1)), int(m.group(2)) def _validate_model_size(model: str, size: str) -> str: if model == "ideogram4": if size not in IDEOGRAM_ALLOWED: return IDEOGRAM_DEFAULT return size if model.startswith("flux"): if size not in FLUX_ALLOWED: return DEFAULT_SIZE return size return DEFAULT_SIZE def _build_ideogram4_workflow(prompt: str, width: int, height: int, seed: int, steps: int = 20, cfg: float = 1.0, negative_prompt: str = "") -> dict: """ Ideogram 4 requires asymmetric CFG wiring: - Main UNet goes through CFGOverride before DualModelGuider - Separate unconditional UNet loaded via second UNETLoader - ConditioningZeroOut for the negative (not a text prompt) - DualModelGuider takes model_negative (unconditional UNet), not model_1 """ return { "1": {"inputs": {"vae_name": "flux2-vae.safetensors"}, "class_type": "VAELoader"}, "2": {"inputs": {"unet_name": "ideogram4_fp8_scaled.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"}, "3": {"inputs": {"clip_name": "qwen3vl_8b_fp8_scaled.safetensors", "type": "ideogram4", "device": "default"}, "class_type": "CLIPLoader"}, "4": {"inputs": {"text": prompt, "clip": ["3", 0]}, "class_type": "CLIPTextEncode"}, "5": {"inputs": {"width": width, "height": height, "batch_size": 1}, "class_type": "EmptyFlux2LatentImage"}, "6": {"inputs": {"noise_seed": seed, "batch_count": 1}, "class_type": "RandomNoise"}, "7": {"inputs": {"sampler_name": "euler"}, "class_type": "KSamplerSelect"}, "8": {"inputs": {"steps": steps, "width": width, "height": height, "mu": 0.5, "std": 1.75}, "class_type": "Ideogram4Scheduler"}, # Asymmetric CFG: second unconditional UNet "15": {"inputs": {"unet_name": "ideogram4_unconditional_fp8_scaled.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"}, # Zero out conditioning for negative path "16": {"inputs": {"conditioning": ["4", 0]}, "class_type": "ConditioningZeroOut"}, # CFG override on main model (70%-100% of steps) "17": {"inputs": {"cfg": 3.0, "start_percent": 0.7, "end_percent": 1.0, "model": ["2", 0]}, "class_type": "CFGOverride"}, # DualModelGuider: main model (CFG overridden), unconditional model, positive, zeroed negative "9": {"inputs": {"cfg": 7.0, "model": ["17", 0], "positive": ["4", 0], "model_negative": ["15", 0], "negative": ["16", 0]}, "class_type": "DualModelGuider"}, "10": {"inputs": {"noise": ["6", 0], "guider": ["9", 0], "sampler": ["7", 0], "sigmas": ["8", 0], "latent_image": ["5", 0]}, "class_type": "SamplerCustomAdvanced"}, "11": {"inputs": {"samples": ["10", 0], "vae": ["1", 0]}, "class_type": "VAEDecode"}, "12": {"inputs": {"filename_prefix": "ideogram4_api", "images": ["11", 0]}, "class_type": "SaveImage"}, } def _build_flux2_workflow(prompt: str, width: int, height: int, seed: int, steps: int = 20, cfg: float = 3.5, negative_prompt: str = "") -> dict: """FLUX.2-klein-9B workflow using UnetLoaderGGUF + CLIPLoaderGGUF (Qwen3VL).""" return { "1": {"inputs": {"vae_name": "flux2-vae.safetensors"}, "class_type": "VAELoader"}, "2": {"inputs": {"unet_name": "flux-2-klein-9b-Q4_K_S.gguf"}, "class_type": "UnetLoaderGGUF"}, "3": {"inputs": {"model": ["2", 0], "max_shift": 1.15, "base_shift": 0.5, "width": width, "height": height}, "class_type": "ModelSamplingFlux"}, # FLUX.2 uses Qwen3VL via CLIPLoaderGGUF, NOT DualCLIPLoader (CLIP-L + T5XXL) "4": {"inputs": {"clip_name": "Qwen3VL-8B-Instruct-Q4_K_M.gguf", "type": "flux2"}, "class_type": "CLIPLoaderGGUF"}, "5": {"inputs": {"text": prompt, "clip": ["4", 0]}, "class_type": "CLIPTextEncode"}, "6": {"inputs": {"conditioning": ["5", 0], "guidance": cfg}, "class_type": "FluxGuidance"}, "7": {"inputs": {"width": width, "height": height, "batch_size": 1}, "class_type": "EmptyFlux2LatentImage"}, "8": {"inputs": {"noise_seed": seed}, "class_type": "RandomNoise"}, "9": {"inputs": {"sampler_name": "euler"}, "class_type": "KSamplerSelect"}, "10": {"inputs": {"model": ["3", 0], "steps": steps, "denoise": 1.0, "scheduler": "simple", "sampler": ["9", 0]}, "class_type": "BasicScheduler"}, "11": {"inputs": {"model": ["3", 0], "conditioning": ["6", 0]}, "class_type": "BasicGuider"}, "12": {"inputs": {"noise": ["8", 0], "guider": ["11", 0], "sampler": ["9", 0], "sigmas": ["10", 0], "latent_image": ["7", 0]}, "class_type": "SamplerCustomAdvanced"}, "13": {"inputs": {"samples": ["12", 0], "vae": ["1", 0]}, "class_type": "VAEDecode"}, "14": {"inputs": {"filename_prefix": "flux2_api", "images": ["13", 0]}, "class_type": "SaveImage"}, } def _build_workflow(model: str, prompt: str, width: int, height: int, seed: int, steps: int, cfg: float, negative_prompt: str = "", style: str = "") -> dict: if model == "ideogram4": return _build_ideogram4_workflow(prompt, width, height, seed, steps, cfg, negative_prompt) elif model == "flux2": return _build_flux2_workflow(prompt, width, height, seed, steps, cfg, negative_prompt) else: raise ValueError(f"Unknown model: {model}") async def _comfy_post(path: str, data: dict, timeout: float = 2400) -> dict: async with httpx.AsyncClient() as client: r = await client.post(f"{COMFY_URL}{path}", json=data, timeout=timeout) if r.status_code >= 400: try: detail = r.json() except Exception: detail = r.text raise HTTPException(status_code=r.status_code, detail=detail) return r.json() async def _comfy_get(path: str, timeout: float = 30) -> dict: async with httpx.AsyncClient() as client: r = await client.get(f"{COMFY_URL}{path}", timeout=timeout) r.raise_for_status() return r.json() async def _generate_image(model: str, prompt: str, width: int, height: int, seed: int, steps: int, cfg: float, negative_prompt: str = "", style: str = "") -> bytes: workflow = _build_workflow(model, prompt, width, height, seed, steps, cfg, negative_prompt, style) resp = await _comfy_post("/prompt", {"prompt": workflow}) prompt_id = resp.get("prompt_id") if not prompt_id: raise HTTPException(status_code=500, detail="No prompt_id from ComfyUI") # Poll for completion - up to 30 min (180 iterations x 10s) for i in range(180): await asyncio.sleep(10) try: hist = await _comfy_get(f"/history/{prompt_id}", timeout=10) except Exception: continue if prompt_id not in hist: continue d = hist[prompt_id] status = d.get("status", {}) status_str = status.get("status_str", "?") if status_str == "error": errors = [] for node_id, msgs in status.get("messages", {}).items(): for level, msg, detail, _ in msgs: if level == "error": errors.append(f"[{node_id}] {msg}: {detail}") detail = d.get("outputs", {}).get("error", "") or "; ".join(errors) or "ComfyUI execution error" raise HTTPException(status_code=500, detail=detail) if status_str == "success": outputs = d.get("outputs", {}) for node_id, node_out in outputs.items(): for img in node_out.get("images", []): fname = img.get("filename") subfolder = img.get("subfolder", "") ftype = img.get("type", "output") url = f"{COMFY_URL}/view?filename={fname}&subfolder={subfolder}&type={ftype}" async with httpx.AsyncClient() as client: img_resp = await client.get(url, timeout=60) img_resp.raise_for_status() return img_resp.content raise HTTPException(status_code=500, detail="No image in ComfyUI outputs") raise HTTPException(status_code=504, detail="ComfyUI generation timeout") app = FastAPI(title="ComfyUI Image Gen OpenAI API", version="2.1.0") class ImageGenerationRequest(BaseModel): model: str = Field(default="ideogram4", description="Model: ideogram4, flux2") prompt: str = Field(..., description="Text prompt") n: int = Field(default=1, ge=1, le=4, description="Number of images") size: str = Field(default="1024x1024") response_format: str = Field(default="url", pattern="^(url|b64_json)$") seed: Optional[int] = Field(default=None, description="Random seed (int). Random if omitted.") steps: Optional[int] = Field(default=None, ge=1, le=100, description="Denoising steps. Overrides quality.") cfg: Optional[float] = Field(default=None, ge=0.0, le=100.0, description="CFG / guidance scale") quality: Optional[str] = Field(default="medium", description="low|medium|high|ultra -> maps to steps") negative_prompt: Optional[str] = Field(default="", description="Negative prompt") style: Optional[str] = Field(default="", description="Style modifier (model-dependent)") user: Optional[str] = Field(default=None, description="OpenAI user field (logged, not used)") class ImageData(BaseModel): url: Optional[str] = None b64_json: Optional[str] = None revised_prompt: Optional[str] = None class ImageGenerationResponse(BaseModel): created: int data: list[ImageData] @app.middleware("http") async def auth_middleware(request: Request, call_next): if API_KEY and request.url.path.startswith("/v1"): auth = request.headers.get("authorization", "") if not auth.startswith("Bearer ") or auth[7:] != API_KEY: return JSONResponse(status_code=401, content={"error": "Unauthorized"}) return await call_next(request) @app.post("/v1/images/generations", response_model=ImageGenerationResponse) async def create_image_generation(req: ImageGenerationRequest): if req.model not in ALLOWED_MODELS: raise HTTPException( status_code=400, detail=f"Invalid model: {req.model}. Allowed: {ALLOWED_MODELS}" ) validated_size = _validate_model_size(req.model, req.size) if validated_size not in ALLOWED_SIZES: raise HTTPException(status_code=400, detail=f"Invalid size: {validated_size}. Allowed: {ALLOWED_SIZES}") width, height = _parse_size(validated_size) created = int(time.time()) data = [] if req.steps is not None: steps = req.steps elif req.quality in QUALITY_STEPS: steps = QUALITY_STEPS[req.quality] else: steps = QUALITY_STEPS["medium"] # Ideogram4: enforce minimum ~12 steps to avoid blank/gray output if req.model == "ideogram4" and steps < 12: steps = 12 if req.model == "ideogram4": cfg = req.cfg if req.cfg is not None else 1.0 elif req.model.startswith("flux"): cfg = req.cfg if req.cfg is not None else 3.5 else: cfg = req.cfg if req.cfg is not None else 7.5 for i in range(req.n): seed = req.seed + i if req.seed is not None else int(time.time() * 1000) + i image_bytes = await _generate_image( req.model, req.prompt, width, height, seed, steps, cfg, req.negative_prompt or "", req.style or "" ) if req.response_format == "b64_json": b64 = base64.b64encode(image_bytes).decode("utf-8") data.append(ImageData(b64_json=b64, revised_prompt=req.prompt)) else: b64 = base64.b64encode(image_bytes).decode("utf-8") data_url = f"data:image/png;base64,{b64}" data.append(ImageData(url=data_url, revised_prompt=req.prompt)) return ImageGenerationResponse(created=created, data=data) @app.get("/v1/models") async def list_models(): """List available image models - dynamically from ALLOWED_MODELS.""" return { "object": "list", "data": [ {"id": m, "object": "model", "owned_by": "comfyui", "permission": []} for m in sorted(ALLOWED_MODELS) ] } @app.get("/health") async def health(): try: await _comfy_get("/system_stats", timeout=5) comfy_status = "ok" except Exception as e: comfy_status = f"error: {e}" return { "status": "ok", "comfyui": COMFY_URL, "comfyui_status": comfy_status, "models": list(ALLOWED_MODELS) } if __name__ == "__main__": import uvicorn port = int(os.environ.get("PORT", "8000")) uvicorn.run(app, host="0.0.0.0", port=port)