95 lines
4.0 KiB
Python
95 lines
4.0 KiB
Python
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
from typing import Optional
|
|
import httpx, json, base64, asyncio, time, random, os
|
|
|
|
app = FastAPI(title="ComfyUI OpenAI Adapter", version="0.1.0")
|
|
|
|
COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://10.0.30.97:8188")
|
|
WORKFLOW_PREFIX = os.environ.get("WORKFLOW_PREFIX", "oai_adapter")
|
|
|
|
# FLUX.2 workflow using Qwen3VL, EmptyFlux2LatentImage, and TAEF2
|
|
# Adapt model names to your local ComfyUI installation.
|
|
FLUX2_WORKFLOW = {
|
|
"1": {"inputs": {"unet_name": "flux-2-klein-9b-Q4_K_S.gguf"}, "class_type": "UnetLoaderGGUF"},
|
|
"2": {"inputs": {"clip_name": "Qwen3VL-8B-Instruct-Q4_K_M.gguf", "type": "flux2"}, "class_type": "CLIPLoaderGGUF"},
|
|
"3": {"inputs": {"vae_name": "taef2"}, "class_type": "VAELoader"},
|
|
"4": {"inputs": {"width": 1024, "height": 1024, "batch_size": 1}, "class_type": "EmptyFlux2LatentImage"},
|
|
"5": {"inputs": {"text": "prompt", "clip": ["2", 0]}, "class_type": "CLIPTextEncode"},
|
|
"6": {"inputs": {"text": "", "clip": ["2", 0]}, "class_type": "CLIPTextEncode"},
|
|
"7": {"inputs": {"seed": 42, "steps": 4, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple",
|
|
"denoise": 1.0, "model": ["1", 0], "positive": ["5", 0], "negative": ["6", 0],
|
|
"latent_image": ["4", 0]}, "class_type": "KSampler"},
|
|
"8": {"inputs": {"samples": ["7", 0], "vae": ["3", 0]}, "class_type": "VAEDecode"},
|
|
"9": {"inputs": {"filename_prefix": WORKFLOW_PREFIX, "images": ["8", 0]}, "class_type": "SaveImage"}
|
|
}
|
|
|
|
class ImageRequest(BaseModel):
|
|
prompt: str
|
|
n: Optional[int] = 1
|
|
size: Optional[str] = "1024x1024"
|
|
response_format: Optional[str] = "b64_json"
|
|
|
|
async def queue_workflow(wf: dict) -> str:
|
|
async with httpx.AsyncClient() as client:
|
|
r = await client.post(f"{COMFYUI_URL}/api/prompt", json={"prompt": wf})
|
|
r.raise_for_status()
|
|
return r.json()["prompt_id"]
|
|
|
|
async def wait_for_image(pid: str, timeout: int = 300) -> list:
|
|
async with httpx.AsyncClient() as client:
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
r = await client.get(f"{COMFYUI_URL}/history/{pid}")
|
|
if r.status_code == 200:
|
|
data = r.json()
|
|
if pid in data and "outputs" in data[pid]:
|
|
imgs = []
|
|
for node_out in data[pid]["outputs"].values():
|
|
if "images" in node_out:
|
|
imgs.extend(node_out["images"])
|
|
if imgs:
|
|
return imgs
|
|
await asyncio.sleep(1)
|
|
raise TimeoutError(f"Timeout waiting for {pid}")
|
|
|
|
async def fetch_image(filename: str, subfolder: str = "", folder_type: str = "output") -> bytes:
|
|
async with httpx.AsyncClient() as client:
|
|
params = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
r = await client.get(f"{COMFYUI_URL}/view", params=params)
|
|
r.raise_for_status()
|
|
return r.content
|
|
|
|
@app.post("/v1/images/generations")
|
|
async def generate(req: ImageRequest):
|
|
imgs = []
|
|
for i in range(req.n):
|
|
wf = json.loads(json.dumps(FLUX2_WORKFLOW))
|
|
wf["5"]["inputs"]["text"] = req.prompt
|
|
w, h = (1024, 1024)
|
|
if "x" in req.size:
|
|
try:
|
|
w, h = map(int, req.size.split("x"))
|
|
except ValueError:
|
|
pass
|
|
wf["4"]["inputs"]["width"] = w
|
|
wf["4"]["inputs"]["height"] = h
|
|
wf["7"]["inputs"]["seed"] = random.randint(1, 2**32)
|
|
|
|
pid = await queue_workflow(wf)
|
|
files = await wait_for_image(pid)
|
|
if not files:
|
|
raise RuntimeError("No images returned")
|
|
data = await fetch_image(files[0]["filename"], files[0].get("subfolder", ""), files[0].get("type", "output"))
|
|
imgs.append({"b64_json": base64.b64encode(data).decode()})
|
|
|
|
return {"created": int(time.time()), "data": imgs}
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "ok"}
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=9000)
|