neon-chat/comfyui_integration.py

"""
NeonChat ComfyUI Integration — Profilbild-Generierung für Characters
"""

import json, os, time, urllib.request, urllib.error
from pathlib import Path

COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://localhost:8188")
COMFYUI_OUTPUT_DIR = Path(os.environ.get("COMFYUI_OUTPUT", str(Path.home() / "comfy" / "output")))
CHARACTERS_DIR = Path(__file__).parent / "characters"
AVATARS_DIR = Path(__file__).parent / "static" / "avatars"

def is_comfyui_running() -> bool:
    try:
        req = urllib.request.Request(f"{COMFYUI_URL}/system_stats", method="GET")
        with urllib.request.urlopen(req, timeout=3) as resp:
            return resp.status == 200
    except:
        return False

def build_portrait_workflow(prompt: str, negative: str = "", width: int = 512, height: int = 512) -> dict:
    """Baut einen ComfyUI Workflow für ein Porträtfoto."""
    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": "noobaiXLNAIFP_vPred10Version.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:
    """Schickt den Workflow an ComfyUI und gibt die Prompt-ID zurück."""
    data = json.dumps({"prompt": workflow}).encode("utf-8")
    req = urllib.request.Request(
        f"{COMFYUI_URL}/prompt",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        result = json.loads(resp.read())
        return result.get("prompt_id", "")

def get_history(prompt_id: str) -> dict:
    """Holt den Verlauf eines Prompts um das fertige Bild zu finden."""
    try:
        req = urllib.request.Request(f"{COMFYUI_URL}/history/{prompt_id}", method="GET")
        with urllib.request.urlopen(req, timeout=5) as resp:
            return json.loads(resp.read())
    except:
        return {}

def wait_for_image(prompt_id: str, timeout: int = 120) -> str | None:
    """Wartet bis das Bild fertig ist und gibt den Dateinamen zurück."""
    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(2)
    return None

def generate_portrait(character_name: str, description: str) -> str | None:
    """
    Generiert ein Profilbild für einen Character via ComfyUI.
    Gibt den Pfad zur gespeicherten Datei zurück, oder None bei Fehler.
    """
    if not is_comfyui_running():
        return None
    
    AVATARS_DIR.mkdir(parents=True, exist_ok=True)
    
    # Build prompt from character description
    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=120)
        if not filename:
            return None
        
        # Find the image in ComfyUI output
        comfy_output = COMFYUI_OUTPUT_DIR / filename
        if not comfy_output.exists():
            return None
        
        # Copy to avatars dir
        safe_name = character_name.lower().replace(" ", "_")
        dest = AVATARS_DIR / f"{safe_name}.png"
        import shutil
        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