338 lines
14 KiB
Python
338 lines
14 KiB
Python
#!/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)
|