arch_agent
b916792082
- Use httpx instead of urllib for ComfyUI API calls - Fix checkpoint name: NoobAI-XL-v1.1.safetensors - Fix output directory: ~/comfy/ComfyUI/output/ - Mara portrait generated successfully (507KB PNG)
130 lines
4.3 KiB
Python
130 lines
4.3 KiB
Python
"""
|
|
NeonChat ComfyUI Integration — Profilbild-Generierung für Characters
|
|
"""
|
|
|
|
import json, os, time, shutil
|
|
from pathlib import Path
|
|
import httpx
|
|
|
|
COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://localhost:8188")
|
|
COMFYUI_OUTPUT_DIR = Path(os.environ.get("COMFYUI_OUTPUT", str(Path.home() / "comfy" / "ComfyUI" / "output")))
|
|
CHARACTERS_DIR = Path(__file__).parent / "characters"
|
|
AVATARS_DIR = Path(__file__).parent / "static" / "avatars"
|
|
|
|
def is_comfyui_running() -> bool:
|
|
try:
|
|
with httpx.Client(timeout=3) as client:
|
|
resp = client.get(f"{COMFYUI_URL}/system_stats")
|
|
return resp.status_code == 200
|
|
except:
|
|
return False
|
|
|
|
def build_portrait_workflow(prompt: str, negative: str = "", width: int = 512, height: int = 512) -> dict:
|
|
return {
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"seed": int(time.time()) % (2**32),
|
|
"steps": 25,
|
|
"cfg": 7.0,
|
|
"sampler_name": "dpmpp_2m",
|
|
"scheduler": "karras",
|
|
"denoise": 1.0,
|
|
"model": ["4", 0],
|
|
"positive": ["6", 0],
|
|
"negative": ["7", 0],
|
|
"latent_image": ["5", 0]
|
|
}
|
|
},
|
|
"4": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "NoobAI-XL-v1.1.safetensors"}
|
|
},
|
|
"5": {
|
|
"class_type": "EmptyLatentImage",
|
|
"inputs": {"width": width, "height": height, "batch_size": 1}
|
|
},
|
|
"6": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {
|
|
"text": f"portrait, {prompt}, detailed face, high quality, masterpiece, best quality, 1girl, solo",
|
|
"clip": ["4", 1]
|
|
}
|
|
},
|
|
"7": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {
|
|
"text": f"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, {negative}",
|
|
"clip": ["4", 1]
|
|
}
|
|
},
|
|
"8": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["3", 0], "vae": ["4", 2]}
|
|
},
|
|
"9": {
|
|
"class_type": "SaveImage",
|
|
"inputs": {"filename_prefix": "neonchat_portrait", "images": ["8", 0]}
|
|
}
|
|
}
|
|
|
|
def queue_prompt(workflow: dict) -> str:
|
|
with httpx.Client(timeout=10) as client:
|
|
resp = client.post(f"{COMFYUI_URL}/prompt", json={"prompt": workflow})
|
|
result = resp.json()
|
|
return result.get("prompt_id", "")
|
|
|
|
def get_history(prompt_id: str) -> dict:
|
|
try:
|
|
with httpx.Client(timeout=5) as client:
|
|
resp = client.get(f"{COMFYUI_URL}/history/{prompt_id}")
|
|
return resp.json()
|
|
except:
|
|
return {}
|
|
|
|
def wait_for_image(prompt_id: str, timeout: int = 180) -> str | None:
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
history = get_history(prompt_id)
|
|
if prompt_id in history:
|
|
outputs = history[prompt_id].get("outputs", {})
|
|
if "9" in outputs:
|
|
images = outputs["9"].get("images", [])
|
|
if images:
|
|
return images[0].get("filename", "")
|
|
time.sleep(3)
|
|
return None
|
|
|
|
def generate_portrait(character_name: str, description: str) -> str | None:
|
|
if not is_comfyui_running():
|
|
return None
|
|
|
|
AVATARS_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
prompt = f"anime style, {description}, simple background, headshot"
|
|
negative = "blurry, deformed, ugly, realistic"
|
|
|
|
workflow = build_portrait_workflow(prompt, negative, width=512, height=512)
|
|
|
|
try:
|
|
prompt_id = queue_prompt(workflow)
|
|
if not prompt_id:
|
|
return None
|
|
|
|
filename = wait_for_image(prompt_id, timeout=180)
|
|
if not filename:
|
|
return None
|
|
|
|
comfy_output = COMFYUI_OUTPUT_DIR / filename
|
|
if not comfy_output.exists():
|
|
return None
|
|
|
|
safe_name = character_name.lower().replace(" ", "_")
|
|
dest = AVATARS_DIR / f"{safe_name}.png"
|
|
shutil.copy2(str(comfy_output), str(dest))
|
|
|
|
return f"/static/avatars/{safe_name}.png"
|
|
except Exception as e:
|
|
print(f"ComfyUI error: {e}")
|
|
return None
|