natiris/bridges/ComfyBridge.py

#!/usr/bin/env python3
"""
ComfyBridge Working – Funktionierende IPAdapter Integration
"""

import json
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
import requests

PATHS = {
    "state": os.path.expanduser("~/natiris/core/natiris_full_state.json"),
    "output_dir": os.path.expanduser("~/natiris/generated"),
    "base_images": os.path.expanduser("~/natiris/assets/base_images"),
}

COMFY_API = "http://localhost:8188"

TRUST_MAP = [
    {"range": [0, 3], "style": "neutral_portrait", "lighting": "professional cold", "ipadapter_weight": 0.4, "distance": "medium"},
    {"range": [4, 7], "style": "personal_context", "lighting": "warm golden", "ipadapter_weight": 0.6, "distance": "medium-close"},
    {"range": [8, 10], "style": "intimate", "lighting": "intimate soft", "ipadapter_weight": 0.8, "distance": "close"},
]


class ComfyBridgeWorking:
    def __init__(self):
        self.output_dir = Path(PATHS["output_dir"])
        self.output_dir.mkdir(exist_ok=True)
        self.client_id = f"natiris_{uuid.uuid4().hex[:6]}"
        
    def check_health(self):
        try:
            resp = requests.get(f"{COMFY_API}/system_stats", timeout=5)
            return {"reachable": True, "version": resp.json()["system"]["comfyui_version"]}
        except Exception as e:
            return {"reachable": False, "error": str(e)}
    
    def get_style_config(self, trust):
        for entry in TRUST_MAP:
            if entry["range"][0] <= trust <= entry["range"][1]:
                return entry
        return TRUST_MAP[1]
    
    def build_prompt(self, trust=7.0, mood=5):
        style = self.get_style_config(trust)
        
        positive = (
            f"portrait of young woman, {style['lighting']}, "
            f"{style['distance']} shot, mood {mood}/10, "
            "beautiful, consistent face, realistic, 8k"
        )
        
        negative = "ugly, deformed, blurry, low quality, extra limbs"
        
        return {
            "positive": positive,
            "negative": negative,
            "style": style["style"],
            "trust": trust,
            "ipadapter_weight": style["ipadapter_weight"],
            "width": 512,
            "height": 768 if trust > 7 else 512,
        }
    
    def build_basic_workflow(self, prompt_data):
        """Einfacher Workflow ohne IPAdapter"""
        seed = int(time.time() * 1000) % 2147483647
        
        return {
            "1": {"inputs": {"text": prompt_data["positive"], "clip": ["4", 1]}, "class_type": "CLIPTextEncode"},
            "2": {"inputs": {"text": prompt_data["negative"], "clip": ["4", 1]}, "class_type": "CLIPTextEncode"},
            "3": {
                "inputs": {
                    "seed": seed,
                    "steps": 25,
                    "cfg": 7.0,
                    "sampler_name": "euler_ancestral",
                    "scheduler": "karras",
                    "denoise": 1.0,
                    "model": ["4", 0],
                    "positive": ["1", 0],
                    "negative": ["2", 0],
                    "latent_image": ["5", 0]
                },
                "class_type": "KSampler"
            },
            "4": {"inputs": {"ckpt_name": "realisticVisionV60B1_v51HyperVAE.safetensors"}, "class_type": "CheckpointLoaderSimple"},
            "5": {"inputs": {"width": prompt_data["width"], "height": prompt_data["height"], "batch_size": 1}, "class_type": "EmptyLatentImage"},
            "6": {"inputs": {"samples": ["3", 0], "vae": ["4", 2]}, "class_type": "VAEDecode"},
            "7": {"inputs": {"filename_prefix": f"natiris_{prompt_data['style']}", "images": ["6", 0]}, "class_type": "SaveImage"},
        }
    
    def submit_and_wait(self, workflow):
        """Sendet Workflow und wartet auf Ergebnis"""
        # Submit
        data = {"prompt": workflow, "client_id": self.client_id}
        resp = requests.post(f"{COMFY_API}/prompt", json=data, timeout=10)
        result = resp.json()
        
        if "prompt_id" not in result:
            return {"success": False, "error": result.get("error", "Submit failed")}
        
        prompt_id = result["prompt_id"]
        print(f"⏳ Generating... (ID: {prompt_id[:8]})")
        
        # Warten (simpler Polling)
        for _ in range(300):  # max 5 Min
            time.sleep(1)
            try:
                history = requests.get(f"{COMFY_API}/history", timeout=5).json()
                if prompt_id in history:
                    return {"success": True, "data": history[prompt_id], "prompt_id": prompt_id}
            except:
                continue
        
        return {"success": False, "error": "Timeout"}
    
    def save_image(self, history_data, prompt_data):
        """Speichert generiertes Bild"""
        outputs = history_data.get("outputs", {})
        
        for node_id, node_out in outputs.items():
            if "images" in node_out:
                for img in node_out["images"]:
                    try:
                        params = {
                            "filename": img["filename"],
                            "subfolder": img.get("subfolder", ""),
                            "type": "output"
                        }
                        resp = requests.get(f"{COMFY_API}/view", params=params, timeout=30)
                        
                        if resp.status_code == 200:
                            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
                            filename = f"natiris_{prompt_data['style']}_{ts}.png"
                            filepath = self.output_dir / filename
                            
                            with open(filepath, "wb") as f:
                                f.write(resp.content)
                            
                            return {
                                "success": True,
                                "path": str(filepath),
                                "filename": filename
                            }
                    except Exception as e:
                        print(f"Save error: {e}")
        
        return {"success": False, "error": "Could not save image"}
    
    def generate(self, state_path=None):
        """Hauptmethode"""
        # State laden für Trust-Level
        trust = 7.0
        mood = 5
        
        if os.path.exists(PATHS["state"]):
            try:
                with open(PATHS["state"]) as f:
                    state = json.load(f)
                    core = state.get("core_state", {})
                    trust = core.get("trust", 7.0)
                    mood = core.get("mood", 5)
            except:
                pass
        
        # Health Check
        health = self.check_health()
        if not health["reachable"]:
            return {"success": False, "error": "ComfyUI unreachable"}
        
        # Prompt & Workflow
        prompt_data = self.build_prompt(trust, mood)
        workflow = self.build_basic_workflow(prompt_data)
        
        # Generieren
        print(f"🎨 Generating with trust={trust}, mood={mood}...")
        result = self.submit_and_wait(workflow)
        
        if not result["success"]:
            return result
        
        # Speichern
        return self.save_image(result["data"], prompt_data)


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--check", action="store_true")
    parser.add_argument("--test", action="store_true")
    args = parser.parse_args()
    
    bridge = ComfyBridgeWorking()
    
    if args.check:
        h = bridge.check_health()
        print(f"ComfyUI: {'✅' if h['reachable'] else '❌'} {h.get('version', 'n/a')}")
        return
    
    if args.test:
        print("Testing generation...")
        result = bridge.generate()
        
        if result["success"]:
            print(f"✅ SUCCESS: {result['path']}")
        else:
            print(f"❌ FAILED: {result.get('error', 'Unknown')}")
        
        with open("/tmp/comfy_result.json", "w") as f:
            json.dump(result, f, indent=2)
    else:
        result = bridge.generate()
        print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()