From c1d4f7e2632bf9f1af39211d4369fe519faed852 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Minette=20Kaunism=C3=A4ki?= Date: Fri, 24 Apr 2026 18:34:18 +0200 Subject: [PATCH] init commit --- .gitignore | 13 ++++ .python-version | 1 + README.md | 196 +++++++++++++++++++++++++++++++++++++++++++++++ __init__.py | 25 ++++++ nodes_edit.py | 131 +++++++++++++++++++++++++++++++ nodes_lora.py | 73 ++++++++++++++++++ nodes_t2i.py | 37 +++++++++ nodes_upscale.py | 57 ++++++++++++++ nodes_video.py | 75 ++++++++++++++++++ pyproject.toml | 10 +++ requirements.txt | 2 + utils.py | 180 +++++++++++++++++++++++++++++++++++++++++++ 12 files changed, 800 insertions(+) create mode 100644 .gitignore create mode 100644 .python-version create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes_edit.py create mode 100644 nodes_lora.py create mode 100644 nodes_t2i.py create mode 100644 nodes_upscale.py create mode 100644 nodes_video.py create mode 100644 pyproject.toml create mode 100644 requirements.txt create mode 100644 utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..17478b5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +# Python-generated files +__pycache__/ +*.py[oc] +build/ +dist/ +wheels/ +*.egg-info + +# Virtual environments +.venv + +# macOS +.DS_Store \ No newline at end of file diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..2c07333 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.11 diff --git a/README.md b/README.md new file mode 100644 index 0000000..278bf6c --- /dev/null +++ b/README.md @@ -0,0 +1,196 @@ +# ComfyUI Pruna API + +ComfyUI custom nodes for running [Pruna API](https://docs.api.pruna.ai) models directly inside ComfyUI workflows. + +## Nodes + +### Pruna Text to Image + +Generate images from text prompts using the Pruna API. + +**Inputs** + +| Input | Type | Description | +|---|---|---| +| `prompt` | string | Text description of the image to generate | +| `model` | dropdown | Pruna API model to use | +| `aspect_ratio` | dropdown | Output image aspect ratio | +| `api_key` | string | Your Pruna API key (or set `PRUNA_API_KEY` env var) | + +**Output:** `IMAGE` — compatible with any ComfyUI node that accepts an image. + +**Supported models:** `p-image` + +**Supported aspect ratios:** `1:1`, `16:9`, `9:16`, `4:3`, `3:4`, `3:2`, `2:3` + +--- + +### Pruna Text to Image LoRA + +Generate images from text prompts with a custom LoRA using the Pruna API. + +> LoRA weights must be a HuggingFace URL in the format `huggingface.co//[/]`. LoRAs must be trained with `p-image-trainer`. + +**Inputs** + +| Input | Type | Description | +|---|---|---| +| `prompt` | string | Text description of the image to generate | +| `model` | dropdown | Pruna API model to use | +| `lora_weights` | string | HuggingFace path to LoRA weights | +| `lora_scale` | float | LoRA influence scale (−1.0 – 3.0, default 0.5) | +| `aspect_ratio` | dropdown | Output image aspect ratio | +| `api_key` | string | Your Pruna API key (or set `PRUNA_API_KEY` env var) | +| `hf_api_token` | string *(optional)* | HuggingFace token for private repos | + +**Output:** `IMAGE` + +**Supported models:** `p-image-lora` + +**Supported aspect ratios:** `1:1`, `16:9`, `9:16`, `4:3`, `3:4`, `3:2`, `2:3` + +--- + +### Pruna Image Edit + +Edit an image with a text instruction using the Pruna API. + +**Inputs** + +| Input | Type | Description | +|---|---|---| +| `image` | IMAGE | Input image to edit | +| `prompt` | string | Edit instruction | +| `model` | dropdown | Pruna API model to use | +| `aspect_ratio` | dropdown | Output aspect ratio (`match_input_image` preserves the original) | +| `api_key` | string | Your Pruna API key (or set `PRUNA_API_KEY` env var) | + +**Output:** `IMAGE` + +**Supported models:** `p-image-edit` + +**Supported aspect ratios:** `match_input_image`, `1:1`, `16:9`, `9:16`, `4:3`, `3:4`, `3:2`, `2:3` + +--- + +### Pruna Image Edit LoRA + +Edit an image with a text instruction and a custom LoRA using the Pruna API. + +> LoRA weights must be a HuggingFace URL in the format `huggingface.co//[/]`. LoRAs must be trained with `p-image-edit-trainer`. + +**Inputs** + +| Input | Type | Description | +|---|---|---| +| `image` | IMAGE | Input image to edit | +| `prompt` | string | Edit instruction | +| `model` | dropdown | Pruna API model to use | +| `lora_weights` | string | HuggingFace path to LoRA weights | +| `lora_scale` | float | LoRA influence scale (−1.0 – 3.0, default 0.5) | +| `aspect_ratio` | dropdown | Output aspect ratio | +| `api_key` | string | Your Pruna API key (or set `PRUNA_API_KEY` env var) | +| `hf_api_token` | string *(optional)* | HuggingFace token for private repos | + +**Output:** `IMAGE` + +**Supported models:** `p-image-edit-lora` + +**Supported aspect ratios:** `match_input_image`, `1:1`, `16:9`, `9:16`, `4:3`, `3:4`, `3:2`, `2:3` + +--- + +### Pruna Image Upscale + +AI-upscale an image to a target megapixel size using the Pruna API. + +**Inputs** + +| Input | Type | Description | +|---|---|---| +| `image` | IMAGE | Input image to upscale | +| `target_mp` | int | Target size in megapixels (1–8, default 4) | +| `enhance_details` | boolean | Apply detail enhancement pass | +| `enhance_realism` | boolean | Apply realism enhancement pass | +| `api_key` | string | Your Pruna API key (or set `PRUNA_API_KEY` env var) | + +**Output:** `IMAGE` + +**Model:** `p-image-upscale` (fixed) + +**Pricing:** $0.005/image for 1–4 MP target · $0.01/image for 5–8 MP target + +--- + +### Pruna Video + +Generate a video from a text prompt using the Pruna API. Optionally connect an image to animate it instead of generating from scratch (image-to-video mode). + +**Inputs** + +| Input | Type | Description | +|---|---|---| +| `prompt` | string | Text description of the video | +| `duration` | int | Video length in seconds (1–20, default 5) | +| `resolution` | dropdown | Output resolution (`720p` or `1080p`) | +| `aspect_ratio` | dropdown | Output aspect ratio | +| `draft` | boolean | Use draft mode (faster, lower cost) | +| `api_key` | string | Your Pruna API key (or set `PRUNA_API_KEY` env var) | +| `image` | IMAGE *(optional)* | Source image to animate (image-to-video mode) | + +**Output:** `STRING` (video URL) — pass to a video preview or download node. + +**Model:** `p-video` (fixed) + +**Supported aspect ratios:** `16:9`, `9:16`, `4:3`, `3:4`, `3:2`, `2:3`, `1:1` + +**Pricing:** + +| Resolution | Draft OFF | Draft ON | +|---|---|---| +| 720p | $0.02/sec | $0.005/sec | +| 1080p | $0.04/sec | $0.01/sec | + +--- + +## Installation + +### Via ComfyUI Manager (recommended) + +Search for **ComfyUI Pruna API** in the ComfyUI Manager and click Install. + +### Manual installation + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/PrunaAI/ComfyUI_pruna_api +pip install -r ComfyUI_pruna_api/requirements.txt +``` + +Restart ComfyUI. The nodes appear under the **pruna ai** category in the node browser. + +## Authentication + +Get your API key from the [Pruna developer portal](https://docs.api.pruna.ai). + +You can provide your key in two ways: + +**Option 1 — Node field:** Type your key directly into the `api_key` input on the node. + +**Option 2 — Environment variable (recommended):** Set `PRUNA_API_KEY` before launching ComfyUI: + +```bash +# macOS / Linux +export PRUNA_API_KEY="your-key-here" +python main.py + +# Windows +set PRUNA_API_KEY="your-key-here" +python main.py +``` + +If both are provided, the node field takes priority. + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..37c240e --- /dev/null +++ b/__init__.py @@ -0,0 +1,25 @@ +from .nodes_t2i import PrunaTextToImage +from .nodes_lora import PrunaTextToImageLoRA +from .nodes_edit import PrunaImageEdit, PrunaImageEditLoRA +from .nodes_upscale import PrunaImageUpscale +from .nodes_video import PrunaVideo + +NODE_CLASS_MAPPINGS = { + "PrunaTextToImage": PrunaTextToImage, + "PrunaTextToImageLoRA": PrunaTextToImageLoRA, + "PrunaImageEdit": PrunaImageEdit, + "PrunaImageEditLoRA": PrunaImageEditLoRA, + "PrunaImageUpscale": PrunaImageUpscale, + "PrunaVideo": PrunaVideo, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PrunaTextToImage": "p-image", + "PrunaTextToImageLoRA": "p-image-lora", + "PrunaImageEdit": "p-image-edit", + "PrunaImageEditLoRA": "p-image-edit-lora", + "PrunaImageUpscale": "p-image-upscale", + "PrunaVideo": "p-video", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes_edit.py b/nodes_edit.py new file mode 100644 index 0000000..807044b --- /dev/null +++ b/nodes_edit.py @@ -0,0 +1,131 @@ +import torch +from .utils import ( + ASPECT_RATIOS, + add_hf_token, + call_api_with_fallback, + generation_url_to_tensor, + resolve_api_key, + upload_image, +) + +EDIT_MODELS = [ + "p-image-edit", +] + +EDIT_ASPECT_RATIOS = ["match_input_image"] + ASPECT_RATIOS + +EDIT_LORA_MODELS = [ + "p-image-edit-lora", +] + + +class PrunaImageEdit: + """Edit an image with a text instruction using the Pruna API. + + Accepts a single ComfyUI IMAGE. The image is uploaded to the Pruna files + API and passed as a URL to the editing model. + """ + + CATEGORY = "pruna ai" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "edit" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True, "default": ""}), + "model": (EDIT_MODELS, {"default": "p-image-edit"}), + "aspect_ratio": (EDIT_ASPECT_RATIOS, {"default": "match_input_image"}), + "api_key": ("STRING", {"default": ""}), + } + } + + def edit( + self, + image: torch.Tensor, + prompt: str, + model: str, + aspect_ratio: str, + api_key: str, + ): + key = resolve_api_key(api_key) + image_url = upload_image(image, key) + payload = { + "input": { + "prompt": prompt, + "images": [image_url], + "aspect_ratio": aspect_ratio, + } + } + data = call_api_with_fallback(model, payload, key) + tensor = generation_url_to_tensor(data) + return (tensor,) + + +class PrunaImageEditLoRA: + """Edit an image with a text instruction and a custom LoRA using the Pruna API. + + LoRA weights must be a HuggingFace URL in the format: + huggingface.co//[/] + + Note: LoRAs used with p-image-edit-lora must be trained with p-image-edit-trainer. + """ + + CATEGORY = "pruna ai" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "edit" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True, "default": ""}), + "model": (EDIT_LORA_MODELS, {"default": "p-image-edit-lora"}), + "lora_weights": ( + "STRING", + {"default": "huggingface.co/username/repo"}, + ), + "lora_scale": ( + "FLOAT", + {"default": 0.5, "min": -1.0, "max": 3.0, "step": 0.05}, + ), + "aspect_ratio": (EDIT_ASPECT_RATIOS, {"default": "match_input_image"}), + "api_key": ("STRING", {"default": ""}), + }, + "optional": { + "hf_api_token": ("STRING", {"default": ""}), + }, + } + + def edit( + self, + image: torch.Tensor, + prompt: str, + model: str, + lora_weights: str, + lora_scale: float, + aspect_ratio: str, + api_key: str, + hf_api_token: str = "", + ): + key = resolve_api_key(api_key) + image_url = upload_image(image, key) + payload: dict = { + "input": { + "prompt": prompt, + "images": [image_url], + "lora_weights": lora_weights, + "lora_scale": lora_scale, + "aspect_ratio": aspect_ratio, + } + } + add_hf_token(payload, hf_api_token) + + data = call_api_with_fallback(model, payload, key) + tensor = generation_url_to_tensor(data) + return (tensor,) diff --git a/nodes_lora.py b/nodes_lora.py new file mode 100644 index 0000000..bd6a993 --- /dev/null +++ b/nodes_lora.py @@ -0,0 +1,73 @@ +from .utils import ( + ASPECT_RATIOS, + add_hf_token, + call_api_with_fallback, + generation_url_to_tensor, + resolve_api_key, +) + +LORA_MODELS = [ + "p-image-lora", +] + + +class PrunaTextToImageLoRA: + """Generate an image from a text prompt with a custom LoRA using the Pruna API. + + LoRA weights must be a HuggingFace URL in the format: + huggingface.co//[/] + + Note: LoRAs used with p-image-lora must be trained with p-image-trainer. + """ + + CATEGORY = "pruna ai" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "generate" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": ""}), + "model": (LORA_MODELS, {"default": "p-image-lora"}), + "lora_weights": ( + "STRING", + {"default": "huggingface.co/username/repo"}, + ), + "lora_scale": ( + "FLOAT", + {"default": 0.5, "min": -1.0, "max": 3.0, "step": 0.05}, + ), + "aspect_ratio": (ASPECT_RATIOS, {"default": "1:1"}), + "api_key": ("STRING", {"default": ""}), + }, + "optional": { + "hf_api_token": ("STRING", {"default": ""}), + }, + } + + def generate( + self, + prompt: str, + model: str, + lora_weights: str, + lora_scale: float, + aspect_ratio: str, + api_key: str, + hf_api_token: str = "", + ): + key = resolve_api_key(api_key) + payload: dict = { + "input": { + "prompt": prompt, + "lora_weights": lora_weights, + "lora_scale": lora_scale, + "aspect_ratio": aspect_ratio, + } + } + add_hf_token(payload, hf_api_token) + + data = call_api_with_fallback(model, payload, key) + tensor = generation_url_to_tensor(data) + return (tensor,) diff --git a/nodes_t2i.py b/nodes_t2i.py new file mode 100644 index 0000000..96d5e7c --- /dev/null +++ b/nodes_t2i.py @@ -0,0 +1,37 @@ +from .utils import ( + ASPECT_RATIOS, + call_api_with_fallback, + generation_url_to_tensor, + resolve_api_key, +) + +T2I_MODELS = [ + "p-image", +] + + +class PrunaTextToImage: + """Generate an image from a text prompt using the Pruna API.""" + + CATEGORY = "pruna ai" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "generate" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": ""}), + "model": (T2I_MODELS, {"default": "p-image"}), + "aspect_ratio": (ASPECT_RATIOS, {"default": "1:1"}), + "api_key": ("STRING", {"default": ""}), + } + } + + def generate(self, prompt: str, model: str, aspect_ratio: str, api_key: str): + key = resolve_api_key(api_key) + payload = {"input": {"prompt": prompt, "aspect_ratio": aspect_ratio}} + data = call_api_with_fallback(model, payload, key) + tensor = generation_url_to_tensor(data) + return (tensor,) diff --git a/nodes_upscale.py b/nodes_upscale.py new file mode 100644 index 0000000..66c59eb --- /dev/null +++ b/nodes_upscale.py @@ -0,0 +1,57 @@ +import torch +from .utils import ( + call_api_with_fallback, + generation_url_to_tensor, + resolve_api_key, + upload_image, +) + + +class PrunaImageUpscale: + """Upscale an image using the Pruna API (p-image-upscale). + + Pricing: $0.005/image for 1-4 MP target, $0.01/image for 5-8 MP target. + """ + + CATEGORY = "pruna ai" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "upscale" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "target_mp": ( + "INT", + {"default": 4, "min": 1, "max": 8, "step": 1}, + ), + "enhance_details": ("BOOLEAN", {"default": False}), + "enhance_realism": ("BOOLEAN", {"default": False}), + "api_key": ("STRING", {"default": ""}), + } + } + + def upscale( + self, + image: torch.Tensor, + target_mp: int, + enhance_details: bool, + enhance_realism: bool, + api_key: str, + ): + key = resolve_api_key(api_key) + image_url = upload_image(image, key) + payload = { + "input": { + "image": image_url, + "target": target_mp, + "enhance_details": enhance_details, + "enhance_realism": enhance_realism, + "output_format": "png", + } + } + data = call_api_with_fallback("p-image-upscale", payload, key) + tensor = generation_url_to_tensor(data) + return (tensor,) diff --git a/nodes_video.py b/nodes_video.py new file mode 100644 index 0000000..be49947 --- /dev/null +++ b/nodes_video.py @@ -0,0 +1,75 @@ +import torch +from .utils import ( + VIDEO_ASPECT_RATIOS, + call_api_with_fallback, + resolve_api_key, + resolve_generation_url, + upload_image, +) + + +def _get_video_url(data: dict) -> str: + url = data.get("generation_url", "") + if not url: + raise RuntimeError(f"Pruna API returned no generation_url. Response: {data}") + return resolve_generation_url(url) + + +class PrunaVideo: + """Generate a video from a text prompt and/or an image using the Pruna p-video model. + + When an image is connected the model animates it; without an image it generates + from the text prompt alone. Returns the video URL as a STRING. + + Pricing (p-video): + 720p draft OFF — $0.02/sec | 720p draft ON — $0.005/sec + 1080p draft OFF — $0.04/sec | 1080p draft ON — $0.01/sec + """ + + CATEGORY = "pruna ai" + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("video_url",) + FUNCTION = "generate" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": ""}), + "duration": ("INT", {"default": 5, "min": 1, "max": 20, "step": 1}), + "resolution": (["720p", "1080p"], {"default": "720p"}), + "aspect_ratio": (VIDEO_ASPECT_RATIOS, {"default": "16:9"}), + "draft": ("BOOLEAN", {"default": False}), + "api_key": ("STRING", {"default": ""}), + }, + "optional": { + "image": ("IMAGE",), + }, + } + + def generate( + self, + prompt: str, + duration: int, + resolution: str, + aspect_ratio: str, + draft: bool, + api_key: str, + image: torch.Tensor = None, + ): + key = resolve_api_key(api_key) + payload: dict = { + "input": { + "prompt": prompt, + "duration": duration, + "resolution": resolution, + "aspect_ratio": aspect_ratio, + "draft": draft, + } + } + + if image is not None: + payload["input"]["image"] = upload_image(image, key) + + data = call_api_with_fallback("p-video", payload, key) + return (_get_video_url(data),) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..d11cffb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,10 @@ +[project] +name = "comfyui-pruna" +version = "0.1.0" +description = "ComfyUI custom nodes for the Pruna API" +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ + "requests>=2.28", + "Pillow>=9.0", +] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..eee536f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +requests>=2.28 +Pillow>=9.0 diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..eb6e8d3 --- /dev/null +++ b/utils.py @@ -0,0 +1,180 @@ +"""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]