hendrik/tools/imagegen.py

import os
import json
import time
import base64
import urllib.request
import urllib.error
import config

schema_generate_image = {
    "type": "function",
    "function": {
        "name": "generate_image",
        "description": "Generate an image from a text prompt using AI image generation. Saves the result to a file.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Text description of the image to generate"
                },
                "output_path": {
                    "type": "string",
                    "description": "Directory or full path to save the image (e.g. ~/Downloads, ~/Downloads/myart.png)"
                },
                "model": {
                    "type": "string",
                    "description": "OpenRouter model slug (optional, uses default from config)"
                },
                "aspect_ratio": {
                    "type": "string",
                    "description": "Aspect ratio: 1:1, 16:9, 9:16, 4:3, 3:4, etc. (optional)"
                },
                "resolution": {
                    "type": "string",
                    "description": "Resolution tier: 512, 1K, 2K, 4K (optional)"
                },
                "input_images": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "List of image paths or URLs to use as reference for image-to-image generation (optional)"
                }
            },
            "required": ["prompt", "output_path"]
        }
    }
}


def generate_image(prompt, output_path, model=None, aspect_ratio=None, resolution=None, input_images=None):
    if not config.IMAGEGEN_API_KEY:
        return "Error: No OpenRouter API key configured. Set imagegen.provider in config.yaml"
    if not config.IMAGEGEN_MODEL and not model:
        return "Error: No image model configured. Set imagegen.default_model in config.yaml or specify model in tool call"

    model = model or config.IMAGEGEN_MODEL

    payload = {"model": model, "prompt": prompt}
    if aspect_ratio:
        payload["aspect_ratio"] = aspect_ratio
    if resolution:
        payload["resolution"] = resolution

    # Build input_references dari paths/URLs
    if input_images:
        references = []
        for img_ref in input_images:
            img_ref = img_ref.strip()
            if not img_ref:
                continue
            if img_ref.startswith("http://") or img_ref.startswith("https://"):
                references.append({"type": "image_url", "image_url": {"url": img_ref}})
            else:
                full = os.path.expanduser(img_ref)
                if not os.path.isfile(full):
                    return f"Error: reference image not found: {full}"
                with open(full, "rb") as f:
                    b64 = base64.b64encode(f.read()).decode("ascii")
                ref_ext = os.path.splitext(full)[1].lower()
                ref_mime = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp"}.get(ref_ext, "image/png")
                references.append({"type": "image_url", "image_url": {"url": f"data:{ref_mime};base64,{b64}"}})
        if references:
            payload["input_references"] = references

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        "https://openrouter.ai/api/v1/images",
        data=data,
        method="POST",
    )
    req.add_header("Content-Type", "application/json")
    req.add_header("Authorization", f"Bearer {config.IMAGEGEN_API_KEY}")

    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            result = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        return f"Error: HTTP {e.code} - {body[:500]}"
    except Exception as e:
        return f"Error: {e}"

    images = result.get("data", [])
    if not images:
        return f"Error: No images in response - {json.dumps(result)[:500]}"

    img = images[0]
    img_bytes = base64.b64decode(img["b64_json"])

    # Detect actual format from magic bytes
    if img_bytes[:8] == b"\x89PNG\r\n\x1a\n":
        ext = ".png"
    elif img_bytes[:3] == b"\xff\xd8\xff":
        ext = ".jpg"
    elif img_bytes[:4] == b"RIFF" and img_bytes[8:12] == b"WEBP":
        ext = ".webp"
    else:
        media_type = img.get("media_type", "image/png")
        ext_map = {"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}
        ext = ext_map.get(media_type, ".png")

    output_path = os.path.expanduser(output_path)
    if os.path.isdir(output_path):
        filename = f"generated_{int(time.time())}{ext}"
        output_path = os.path.join(output_path, filename)
    elif not os.path.splitext(output_path)[1]:
        output_path += ext

    parent = os.path.dirname(output_path)
    if parent:
        os.makedirs(parent, exist_ok=True)

    with open(output_path, "wb") as f:
        f.write(img_bytes)

    cost = result.get("usage", {}).get("cost", "?")
    return f"Image saved to {output_path} ({len(img_bytes)} bytes, cost: ${cost})"