utils.py

"""Shared helpers for Pruna API nodes."""

import os
import time
from io import BytesIO

import numpy as np
import requests
import torch
from PIL import Image

PRUNA_API_BASE = "https://api.pruna.ai"
PREDICTIONS_URL = f"{PRUNA_API_BASE}/v1/predictions"
FILES_URL = f"{PRUNA_API_BASE}/v1/files"

ASPECT_RATIOS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
VIDEO_ASPECT_RATIOS = ["16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "1:1"]

POLL_INTERVAL = 3   # seconds between status checks
POLL_TIMEOUT = 300  # seconds before giving up


def resolve_api_key(api_key: str) -> str:
    """Return the api_key field value, falling back to the PRUNA_API_KEY env var."""
    key = api_key.strip() if api_key else ""
    if not key:
        key = os.environ.get("PRUNA_API_KEY", "")
    if not key:
        raise RuntimeError(
            "No Pruna API key provided. Set the api_key input or the "
            "PRUNA_API_KEY environment variable."
        )
    return key


def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    """Convert a ComfyUI IMAGE tensor [1, H, W, C] or [H, W, C] to a PIL Image."""
    if tensor.ndim == 4:
        tensor = tensor[0]
    arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
    if arr.shape[-1] == 4:
        return Image.fromarray(arr, mode="RGBA").convert("RGB")
    return Image.fromarray(arr, mode="RGB")


def upload_image(tensor: torch.Tensor, api_key: str) -> str:
    """Upload a ComfyUI IMAGE tensor to the Pruna files API and return the file URL."""
    pil_image = tensor_to_pil(tensor)
    buf = BytesIO()
    pil_image.save(buf, format="PNG")
    buf.seek(0)

    response = requests.post(
        FILES_URL,
        headers={"apikey": api_key},
        files={"content": ("image.png", buf, "image/png")},
        timeout=60,
    )
    if response.status_code not in (200, 201):
        raise RuntimeError(
            f"Pruna file upload failed {response.status_code}: {response.text}"
        )
    data = response.json()
    return data["urls"]["get"]


def call_api_sync(model: str, payload: dict, api_key: str) -> dict:
    """POST a prediction and return the JSON response (synchronous, Try-Sync: true)."""
    headers = {
        "Content-Type": "application/json",
        "apikey": api_key,
        "Model": model,
        "Try-Sync": "true",
    }
    response = requests.post(PREDICTIONS_URL, json=payload, headers=headers, timeout=120)
    if response.status_code not in (200, 201):
        raise RuntimeError(
            f"Pruna API error {response.status_code}: {response.text}"
        )
    return response.json()


def _poll_until_complete(get_url: str, api_key: str) -> dict:
    """Poll a submitted Pruna job by its get_url until it completes or times out."""
    start = time.time()
    while time.time() - start < POLL_TIMEOUT:
        time.sleep(POLL_INTERVAL)
        status_response = requests.get(
            get_url,
            headers={"apikey": api_key},
            timeout=30,
        )
        if status_response.status_code != 200:
            raise RuntimeError(
                f"Pruna status check failed {status_response.status_code}: "
                f"{status_response.text}"
            )
        status = status_response.json()
        state = status.get("status", "")
        if state == "succeeded" or status.get("generation_url"):
            return status
        if state in ("failed", "canceled"):
            raise RuntimeError(f"Pruna job {state}: {status}")

    raise RuntimeError(
        f"Pruna job did not complete within {POLL_TIMEOUT}s. "
        f"Check the Pruna developer portal for job status."
    )


def call_api_async(model: str, payload: dict, api_key: str) -> dict:
    """POST a prediction and poll until complete. Returns the final status dict."""
    headers = {
        "Content-Type": "application/json",
        "apikey": api_key,
        "Model": model,
    }
    response = requests.post(PREDICTIONS_URL, json=payload, headers=headers, timeout=60)
    if response.status_code not in (200, 201):
        raise RuntimeError(
            f"Pruna API error {response.status_code}: {response.text}"
        )
    job = response.json()

    if job.get("status") == "succeeded" or job.get("generation_url"):
        return job

    get_url = job.get("get_url")
    if not get_url:
        raise RuntimeError(f"Pruna API returned no get_url for polling. Response: {job}")

    return _poll_until_complete(get_url, api_key)


def call_api_with_fallback(model: str, payload: dict, api_key: str) -> dict:
    """Try sync prediction first; if no generation_url yet, poll the existing job."""
    data = call_api_sync(model, payload, api_key)
    if data.get("generation_url"):
        return data

    get_url = data.get("get_url")
    if not get_url:
        raise RuntimeError(
            f"Pruna API returned no generation_url and no get_url. Response: {data}"
        )
    return _poll_until_complete(get_url, api_key)


def add_hf_token(payload: dict, hf_api_token: str) -> None:
    """Attach a non-empty HuggingFace token to payload['input']."""
    token = hf_api_token.strip()
    if token:
        payload["input"]["hf_api_token"] = token


def resolve_generation_url(url: str) -> str:
    """Normalise a generation_url that may be a relative path."""
    if url.startswith("/"):
        return f"{PRUNA_API_BASE}{url}"
    return url


def generation_url_to_tensor(data: dict) -> torch.Tensor:
    """Download generation_url from a Pruna API response and return an IMAGE tensor."""
    generation_url = data.get("generation_url")
    if not generation_url:
        raise RuntimeError(f"Pruna API returned no generation_url. Response: {data}")

    generation_url = resolve_generation_url(generation_url)

    img_response = requests.get(generation_url, timeout=60)
    if img_response.status_code != 200:
        raise RuntimeError(
            f"Failed to download image from {generation_url}: "
            f"{img_response.status_code}"
        )

    image = Image.open(BytesIO(img_response.content)).convert("RGB")
    image_np = np.array(image).astype(np.float32) / 255.0
    return torch.from_numpy(image_np).unsqueeze(0)  # [1, H, W, C]