neon-chat/comfyui_integration.py

"""
NeonChat ComfyUI Integration — Profilbild-Generierung für Characters
Basiert auf dem Character-Aussehen (appearance)
"""

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")))
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_prompt_from_appearance(character: dict, context: str = "") -> str:
    """Baut einen ComfyUI-Prompt aus dem Character-Aussehen."""
    app = character.get("appearance", {})
    
    style = app.get("style", "anime")
    gender = app.get("gender", "female")
    age = app.get("age", "25")
    height = app.get("height", "")
    build = app.get("build", "")
    hair_color = app.get("hair_color", "")
    hair_style = app.get("hair_style", "")
    eye_color = app.get("eye_color", "")
    skin = app.get("skin", "")
    clothing = app.get("clothing", "")
    distinctive = app.get("distinctive", "")
    
    # Style prefix
    style_map = {
        "anime": "anime style, anime art, ",
        "realistic": "photorealistic, realistic, ",
        "comic": "comic book style, western comic, ",
        "cartoon": "cartoon style, ",
        "pixel": "pixel art, ",
    }
    style_prefix = style_map.get(style, "anime style, ")
    
    # Gender
    if gender == "female":
        gender_tag = "1girl, woman, "
    elif gender == "male":
        gender_tag = "1boy, man, "
    else:
        gender_tag = "person, "
    
    # Build prompt
    parts = [
        f"portrait, headshot, {style_prefix}{gender_tag}",
        f"{age} years old",
    ]
    if build:
        parts.append(build)
    if hair_color and hair_style:
        parts.append(f"{hair_color} hair, {hair_style}")
    elif hair_color:
        parts.append(f"{hair_color} hair")
    elif hair_style:
        parts.append(hair_style)
    if eye_color:
        parts.append(f"{eye_color} eyes")
    if skin:
        parts.append(f"{skin} skin")
    if clothing:
        parts.append(f"wearing {clothing}")
    if distinctive:
        parts.append(distinctive)
    if context:
        parts.append(context)
    
    prompt = ", ".join(parts)
    prompt += ", detailed face, high quality, masterpiece, best quality, simple background"
    
    return prompt

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": prompt,
                "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, blurry, deformed, ugly, {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, character: dict, context: str = "") -> str | None:
    """Generiert ein Profilbild basierend auf dem Character-Aussehen."""
    if not is_comfyui_running():
        return None
    
    AVATARS_DIR.mkdir(parents=True, exist_ok=True)
    
    prompt = build_prompt_from_appearance(character, context)
    negative = "realistic, 3d, render"
    
    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