hendrik/tools/imagegen.py

import os
import json
import time
import base64
import urllib.request
import urllib.error
from lib import moop

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 _images_url(base_url):
    base_url = (base_url or "").rstrip("/")
    if base_url.endswith("/images"):
        return base_url
    return base_url + "/images"


def _post_images(url, api_key, payload):
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, method="POST")
    req.add_header("Content-Type", "application/json")
    req.add_header("Authorization", f"Bearer {api_key}")
    with urllib.request.urlopen(req, timeout=120) as resp:
        return json.loads(resp.read().decode("utf-8"))


def generate_image(prompt, output_path, model=None, aspect_ratio=None, resolution=None, input_images=None):
    chain = moop.default_chain(None, "imagegen")
    if not chain:
        return "Error: No imagegen model set configured. Atur via Model > Manage Sets (Model Option)."

    candidates = chain
    if model:
        candidates = [c for c in chain if c["model"] == model] or [chain[0]]

    last_err = ""
    for cand in candidates:
        url = _images_url(cand["base_url"])
        key = cand.get("api_key") or ""

        payload = {"model": cand["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

        try:
            result = _post_images(url, key, payload)
            images = result.get("data", [])
            if not images:
                last_err = f"Error: No images in response - {json.dumps(result)[:500]}"
                continue
            return _save_image(result, images[0], output_path)
        except urllib.error.HTTPError as e:
            body = e.read().decode("utf-8", errors="replace")
            last_err = f"Error: HTTP {e.code} - {body[:500]}"
        except Exception as e:
            last_err = f"Error: {e}"

    return last_err


def _save_image(result, img, output_path):
    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})"