scripts/comfyui-client.py

"""
comfyui-client.py — ComfyUI API Client for auto B-roll image generation
Connects to local ComfyUI Desktop and generates images from prompts.

Usage:
  python comfyui-client.py --prompt "cinematic sunset" --output image.png
  python comfyui-client.py --prompt-file prompts.json --output-dir images/
"""
import argparse, json, os, sys, time, urllib.request, urllib.parse, uuid, websocket, ssl

COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://127.0.0.1:8188")

# DreamShaper 8 workflow template
WORKFLOW_TEMPLATE = {
    "3": {
        "class_type": "KSampler",
        "inputs": {
            "cfg": 7.0,
            "denoise": 1.0,
            "latent_image": ["5", 0],
            "model": ["4", 0],
            "negative": ["7", 0],
            "positive": ["6", 0],
            "sampler_name": "euler_ancestral",
            "scheduler": "normal",
            "seed": 42,
            "steps": 25
        }
    },
    "4": {
        "class_type": "CheckpointLoaderSimple",
        "inputs": {
            "ckpt_name": "dreamshaper_8_pruned.safetensors"
        }
    },
    "5": {
        "class_type": "EmptyLatentImage",
        "inputs": {
            "batch_size": 1,
            "height": 512,
            "width": 768
        }
    },
    "6": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": ["4", 1],
            "text": "cinematic photo, beautiful scene, high quality, detailed, 4k"
        }
    },
    "7": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "clip": ["4", 1],
            "text": "blurry, low quality, distorted, ugly, watermark, text"
        }
    },
    "8": {
        "class_type": "VAEDecode",
        "inputs": {
            "samples": ["3", 0],
            "vae": ["4", 2]
        }
    },
    "9": {
        "class_type": "SaveImage",
        "inputs": {
            "filename_prefix": "broll",
            "images": ["8", 0]
        }
    }
}


def check_comfyui():
    """Check if ComfyUI is running."""
    try:
        req = urllib.request.Request(f"{COMFYUI_URL}/system_stats")
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read())
            print(f"  ComfyUI connected: {data.get('system', {}).get('python_version', 'unknown')}")
            return True
    except Exception as e:
        print(f"  ComfyUI not reachable: {e}")
        return False


def queue_prompt(workflow, client_id):
    """Queue a workflow for execution."""
    data = json.dumps({"prompt": workflow, "client_id": client_id}).encode('utf-8')
    req = urllib.request.Request(
        f"{COMFYUI_URL}/prompt",
        data=data,
        headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())


def get_history(prompt_id):
    """Get execution history."""
    req = urllib.request.Request(f"{COMFYUI_URL}/history/{prompt_id}")
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())


def get_image(filename, subfolder="", folder_type="output"):
    """Download generated image."""
    params = urllib.parse.urlencode({"filename": filename, "subfolder": subfolder, "type": folder_type})
    req = urllib.request.Request(f"{COMFYUI_URL}/view?{params}")
    with urllib.request.urlopen(req) as resp:
        return resp.read()


def wait_for_completion(prompt_id, timeout=120):
    """Wait for workflow to complete."""
    start = time.time()
    while time.time() - start < timeout:
        history = get_history(prompt_id)
        if prompt_id in history:
            outputs = history[prompt_id].get("outputs", {})
            for node_id, output in outputs.items():
                if "images" in output:
                    return output["images"]
        time.sleep(1)
    return None


def generate_image(prompt, output_path, seed=None, steps=25, width=768, height=512):
    """Generate a single image from prompt."""
    import copy
    workflow = copy.deepcopy(WORKFLOW_TEMPLATE)

    # Set prompt
    workflow["6"]["inputs"]["text"] = prompt
    # Set seed
    if seed is None:
        seed = int.from_bytes(os.urandom(4), 'big') % (2**32)
    workflow["3"]["inputs"]["seed"] = seed
    workflow["3"]["inputs"]["steps"] = steps
    workflow["5"]["inputs"]["width"] = width
    workflow["5"]["inputs"]["height"] = height

    client_id = str(uuid.uuid4())

    try:
        result = queue_prompt(workflow, client_id)
        prompt_id = result["prompt_id"]
        print(f"    Queued: {prompt_id}")

        images = wait_for_completion(prompt_id, timeout=180)
        if images:
            img_data = get_image(images[0]["filename"], images[0].get("subfolder", ""))
            os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
            with open(output_path, 'wb') as f:
                f.write(img_data)
            print(f"    Saved: {output_path} ({len(img_data)/1024:.0f} KB)")
            return True
        else:
            print(f"    TIMEOUT waiting for image")
            return False
    except Exception as e:
        print(f"    ERROR: {e}")
        return False


def generate_batch(prompts, output_dir, prefix="broll"):
    """Generate multiple images from a list of prompts."""
    os.makedirs(output_dir, exist_ok=True)
    results = []
    for i, prompt in enumerate(prompts):
        output_path = os.path.join(output_dir, f"{prefix}_{i:03d}.png")
        print(f"\n  [{i+1}/{len(prompts)}] {prompt[:60]}...")
        success = generate_image(prompt, output_path, seed=42 + i)
        results.append({"prompt": prompt, "output": output_path, "success": success})
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ComfyUI API Client")
    parser.add_argument("--prompt", help="Single prompt to generate")
    parser.add_argument("--prompt-file", help="JSON file with list of prompts")
    parser.add_argument("--output", default="output.png", help="Output file (single prompt)")
    parser.add_argument("--output-dir", default="images", help="Output directory (batch)")
    parser.add_argument("--width", type=int, default=768)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--steps", type=int, default=25)
    parser.add_argument("--url", default=None, help="ComfyUI URL")
    args = parser.parse_args()

    if args.url:
        COMFYUI_URL = args.url

    print("[1/3] Checking ComfyUI...")
    if not check_comfyui():
        print("\n  ComfyUI is not running!")
        print("  Please start ComfyUI Desktop first.")
        print("  Or set COMFYUI_URL environment variable.")
        sys.exit(1)

    if args.prompt:
        print(f"\n[2/3] Generating: {args.prompt[:60]}...")
        generate_image(args.prompt, args.output, steps=args.steps,
                      width=args.width, height=args.height)
    elif args.prompt_file:
        print(f"\n[2/3] Loading prompts from {args.prompt_file}...")
        with open(args.prompt_file, encoding='utf-8') as f:
            data = json.load(f)
        prompts = [s.get("image_prompt", "") for s in data.get("scenes", data) if s.get("image_prompt")]
        print(f"  Found {len(prompts)} prompts")
        print(f"\n[3/3] Generating {len(prompts)} images...")
        results = generate_batch(prompts, args.output_dir)
        success = sum(1 for r in results if r["success"])
        print(f"\n  Done: {success}/{len(prompts)} images generated")
    else:
        parser.print_help()