BoosAPI/boosapi_comfy_nodes.py

"""
BoosAPI Custom Nodes for ComfyUI
==================================
Call BoosAPI/metapi endpoints directly from ComfyUI workflows:
  - Text generation (gpt-5.5)
  - Image generation (gpt-image-2)
  - Text-to-Speech
  - Video generation (async)

Configuration (in ComfyUI settings or env vars):
  BOOSAPI_BASE  = https://api.boos.lat/v1
  BOOSAPI_KEY   = sk-boos4721
"""

import os
import json
import time
import requests
import torch
import numpy as np
from PIL import Image
from io import BytesIO
import folder_paths

# ── Config ──────────────────────────────────────────────────────────────────

BOOSAPI_BASE = os.getenv("BOOSAPI_BASE", "https://api.boos.lat/v1")
BOOSAPI_KEY = os.getenv("BOOSAPI_KEY", "sk-boos4721")

HEADERS = {
    "Authorization": f"Bearer {BOOSAPI_KEY}",
    "Content-Type": "application/json",
}


def _call_api(endpoint: str, payload: dict, timeout: int = 120) -> dict:
    url = f"{BOOSAPI_BASE}/{endpoint.lstrip('/')}"
    resp = requests.post(url, headers=HEADERS, json=payload, timeout=timeout)
    if not resp.ok:
        raise RuntimeError(f"BoosAPI {endpoint} failed HTTP {resp.status_code}: {resp.text}")
    return resp.json()


# ── Shared helpers ──────────────────────────────────────────────────────────

def _pil_to_tensor(img: Image.Image) -> torch.Tensor:
    """PIL -> ComfyUI IMAGE (1, H, W, 3) float32 0-1."""
    img = img.convert("RGB")
    arr = np.array(img).astype(np.float32) / 255.0
    return torch.from_numpy(arr).unsqueeze(0)


def _tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    """ComfyUI IMAGE -> PIL."""
    arr = (tensor.squeeze(0).cpu().numpy() * 255).astype(np.uint8)
    return Image.fromarray(arr)


# ── Node: Text Generation (gpt-5.5) ─────────────────────────────────────────

class BoosAPITextGen:
    """Call gpt-5.5 chat completion. Useful for prompt engineering within workflows."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": ("STRING", {"default": "gpt-5.5"}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0, "max": 2, "step": 0.1}),
                "max_tokens": ("INT", {"default": 2048, "min": 1, "max": 32768}),
            },
            "optional": {
                "system_prompt": ("STRING", {"multiline": True, "default": ""}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate"
    CATEGORY = "BoosAPI"

    def generate(self, prompt, model, temperature, max_tokens, system_prompt=""):
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        data = _call_api("chat/completions", {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        })
        text = data["choices"][0]["message"]["content"]
        return (text,)


# ── Node: Image Generation (gpt-image-2) ────────────────────────────────────

class BoosAPIImageGen:
    """Generate image via BoosAPI. Replaces CheckpointLoader+KSampler+VAE chain."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": (["gpt-image-2", "dall-e-3"], {"default": "gpt-image-2"}),
                "size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
            },
            "optional": {
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "n": ("INT", {"default": 1, "min": 1, "max": 4}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "image_url")
    FUNCTION = "generate"
    CATEGORY = "BoosAPI"

    def generate(self, prompt, model, size, negative_prompt="", n=1):
        payload = {
            "model": model,
            "prompt": prompt,
            "n": n,
            "size": size,
        }
        if negative_prompt:
            payload["negative_prompt"] = negative_prompt

        data = _call_api("images/generations", payload)

        # Download first image
        image_url = data["data"][0]["url"]
        img_resp = requests.get(image_url, timeout=60)
        img = Image.open(BytesIO(img_resp.content))
        tensor = _pil_to_tensor(img)
        return (tensor, image_url)


# ── Node: Text-to-Speech ────────────────────────────────────────────────────

class BoosAPITTS:
    """Generate speech audio from text via BoosAPI TTS."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True, "default": ""}),
                "voice": (["alloy", "echo", "fable", "onyx", "nova", "shimmer"], {"default": "alloy"}),
            },
            "optional": {
                "speed": ("FLOAT", {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.25}),
            },
        }

    RETURN_TYPES = ("STRING",)  # file path to generated audio
    RETURN_NAMES = ("audio_path",)
    FUNCTION = "generate"
    CATEGORY = "BoosAPI"
    OUTPUT_NODE = True

    def generate(self, text, voice, speed=1.0):
        url = f"{BOOSAPI_BASE}/audio/speech"
        resp = requests.post(
            url, headers=HEADERS,
            json={"model": "tts-1", "input": text, "voice": voice, "speed": speed},
            timeout=120,
        )
        if not resp.ok:
            raise RuntimeError(f"TTS failed HTTP {resp.status_code}: {resp.text}")

        # Save to ComfyUI output directory
        output_dir = folder_paths.get_output_directory()
        os.makedirs(output_dir, exist_ok=True)
        filename = f"boosapi_tts_{int(time.time())}.mp3"
        filepath = os.path.join(output_dir, filename)
        with open(filepath, "wb") as f:
            f.write(resp.content)

        print(f"[BoosAPI] TTS saved: {filepath}")
        return (filepath,)


# ── Node: Video Generation (async) ─────────────────────────────────────────

class BoosAPIVideoGen:
    """Generate video via BoosAPI (async). Returns task ID for polling."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": ("STRING", {"default": "cogvideo-5b"}),
                "size": (["1024x1024", "1024x576", "576x1024", "1920x1080"], {"default": "1024x576"}),
            },
            "optional": {
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "duration": ("INT", {"default": 5, "min": 2, "max": 30}),
                "fps": ("INT", {"default": 16, "min": 8, "max": 30}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("task_id", "status_url")
    FUNCTION = "generate"
    CATEGORY = "BoosAPI"

    def generate(self, prompt, model, size, negative_prompt="", duration=5, fps=16):
        width, height = size.split("x")
        payload = {
            "model": model,
            "prompt": prompt,
            "width": int(width),
            "height": int(height),
            "duration": duration,
            "fps": fps,
        }
        if negative_prompt:
            payload["negative_prompt"] = negative_prompt

        data = _call_api("videos", payload)
        task_id = data.get("id", "")
        status_url = f"{BOOSAPI_BASE}/videos/{task_id}"
        return (task_id, status_url)


class BoosAPIVideoStatus:
    """Poll video generation status. Returns download URL when done."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "task_id": ("STRING", {"default": ""}),
            },
            "optional": {
                "wait": ("BOOLEAN", {"default": True, "label_on": "Wait until done", "label_off": "Check once"}),
                "poll_interval": ("INT", {"default": 5, "min": 1, "max": 60}),
                "max_polls": ("INT", {"default": 60, "min": 1, "max": 600}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("status", "video_url")
    FUNCTION = "poll"
    CATEGORY = "BoosAPI"

    def poll(self, task_id, wait=True, poll_interval=5, max_polls=60):
        url = f"{BOOSAPI_BASE}/videos/{task_id}"

        for attempt in range(max_polls):
            resp = requests.get(url, headers=HEADERS, timeout=30)
            if not resp.ok:
                return ("error", f"HTTP {resp.status_code}")

            data = resp.json()
            status = data.get("status", "unknown")

            if status == "completed":
                video_url = data.get("video_url") or data.get("output", {}).get("video_url", "")
                return ("completed", video_url)
            elif status in ("failed", "error"):
                error_msg = data.get("error", str(data))
                return (f"failed: {error_msg}", "")

            if not wait:
                return (status, "")

            time.sleep(poll_interval)

        return ("timeout", "")


# ── Registration ────────────────────────────────────────────────────────────

NODE_CLASS_MAPPINGS = {
    "BoosAPITextGen": BoosAPITextGen,
    "BoosAPIImageGen": BoosAPIImageGen,
    "BoosAPITTS": BoosAPITTS,
    "BoosAPIVideoGen": BoosAPIVideoGen,
    "BoosAPIVideoStatus": BoosAPIVideoStatus,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BoosAPITextGen": "BoosAPI Text Gen (gpt-5.5)",
    "BoosAPIImageGen": "BoosAPI Image Gen (gpt-image-2)",
    "BoosAPITTS": "BoosAPI Text-to-Speech",
    "BoosAPIVideoGen": "BoosAPI Video Gen (async)",
    "BoosAPIVideoStatus": "BoosAPI Video Status",
}