import base64 import io import time import numpy as np import requests import torch from PIL import Image VIDEO_GENERATIONS_ENDPOINT = "https://api.aimlapi.com/v2/video/generations" VIDEO_SUCCESS_STATES = {"completed", "succeeded", "success", "ready", "finished"} VIDEO_FAILURE_STATES = {"failed", "error", "canceled", "cancelled", "expired"} DEFAULT_VIDEO_TIMEOUT = 480 DEFAULT_VIDEO_POLL_INTERVAL = 5.0 def aspect_ratio_to_size(aspect_ratio: str, resolution: str = "1K") -> str: """Convert aspect ratio to pixel size for Seedream with resolution scaling.""" size_map_1k = { "1:1": "1024x1024", "16:9": "1920x1080", "9:16": "1080x1920", "4:3": "1024x768", "3:4": "768x1024", "3:2": "1152x768", "2:3": "768x1152", "21:9": "2560x1080", "5:4": "1280x1024", "4:5": "1024x1280", } size_map_2k = { "1:1": "2048x2048", "16:9": "2560x1440", "9:16": "1440x2560", "4:3": "2048x1536", "3:4": "1536x2048", "3:2": "2304x1536", "2:3": "1536x2304", "21:9": "3440x1440", "5:4": "2048x1638", "4:5": "1638x2048", } size_map_4k = { "1:1": "4096x4096", "16:9": "3840x2160", "9:16": "2160x3840", "4:3": "4096x3072", "3:4": "3072x4096", "3:2": "4608x3072", "2:3": "3072x4608", "21:9": "5120x2160", "5:4": "4096x3276", "4:5": "3276x4096", } if resolution == "2K": return size_map_2k.get(aspect_ratio, "2048x2048") if resolution == "4K": return size_map_4k.get(aspect_ratio, "4096x4096") return size_map_1k.get(aspect_ratio, "1024x1024") def tensor_to_data_uri(image_tensor: torch.Tensor) -> str: """Serialize a single image tensor into a PNG data URI.""" if not isinstance(image_tensor, torch.Tensor): raise ValueError("Image input must be a torch tensor.") tensor = image_tensor if tensor.dim() == 4: tensor = tensor[0] array = tensor.detach().cpu().numpy() if array.ndim != 3: raise ValueError("Expected image tensor with shape (H, W, C) or (C, H, W).") if array.shape[0] in (1, 3) and array.shape[-1] not in (1, 3): array = np.transpose(array, (1, 2, 0)) array = np.clip(array * 255.0, 0, 255).astype(np.uint8) image = Image.fromarray(array) buffer = io.BytesIO() image.save(buffer, format="PNG") encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:image/png;base64,{encoded}" def tensors_to_data_uri_list(image_batch: torch.Tensor, limit: int = None) -> list: """Convert a batch of tensors to data URIs.""" if image_batch is None: return [] if image_batch.dim() == 3: tensors = [image_batch] elif image_batch.dim() == 4: tensors = [image_batch[i] for i in range(image_batch.shape[0])] else: raise ValueError("Expected image batch tensor with 3 or 4 dimensions.") if limit is not None: tensors = tensors[:limit] return [tensor_to_data_uri(t) for t in tensors] def prepare_image_value(image_batch, image_url, required=False, field_label="image"): """Return a URL or data URI for the provided tensor/URL combo.""" if image_batch is not None: return tensor_to_data_uri(image_batch) if image_url and image_url.strip(): return image_url.strip() if required: raise ValueError(f"{field_label} is required. Supply an image tensor or a URL.") return None def validate_api_key(api_key: str) -> str: key = (api_key or "").strip() if not key: raise ValueError("API key is required. Get one from https://aimlapi.com/app/keys") return key def build_auth_headers(api_key: str) -> dict: key = validate_api_key(api_key) return { "Authorization": f"Bearer {key}", "Content-Type": "application/json", } def raise_video_api_error(response): message = f"Video API request failed with status {response.status_code}" try: error_data = response.json() detail = error_data.get("error") or error_data.get("message") if detail: message += f": {detail}" else: message += f": {error_data}" except Exception: message += f": {response.text[:400]}" raise ValueError(message) def submit_video_generation(api_key: str, payload: dict) -> str: headers = build_auth_headers(api_key) try: response = requests.post(VIDEO_GENERATIONS_ENDPOINT, headers=headers, json=payload, timeout=120) except requests.exceptions.RequestException as exc: raise ValueError(f"Network error while submitting video generation: {exc}") if response.status_code != 200: raise_video_api_error(response) try: data = response.json() except Exception as exc: raise ValueError(f"Unable to parse video generation response: {exc}") generation_id = ( data.get("generation_id") or data.get("id") or data.get("data", {}).get("generation_id") or data.get("data", {}).get("id") ) if not generation_id: raise ValueError(f"Video generation request succeeded but no generation_id was returned: {data}") return generation_id def extract_video_url(payload): if isinstance(payload, dict): for key in ("video_url", "url", "output_url"): value = payload.get(key) if isinstance(value, str) and value.startswith("http"): return value for key in ("result", "results", "output", "outputs", "videos", "data"): value = payload.get(key) if isinstance(value, dict): found = extract_video_url(value) if found: return found elif isinstance(value, list): for item in value: found = extract_video_url(item) if found: return found elif isinstance(payload, list): for item in payload: found = extract_video_url(item) if found: return found return None def poll_video_generation(api_key: str, generation_id: str, timeout: int = DEFAULT_VIDEO_TIMEOUT, interval: float = DEFAULT_VIDEO_POLL_INTERVAL) -> str: headers = build_auth_headers(api_key) params = {"generation_id": generation_id} start_time = time.time() last_status = "queued" while time.time() - start_time < timeout: try: response = requests.get(VIDEO_GENERATIONS_ENDPOINT, headers=headers, params=params, timeout=60) except requests.exceptions.RequestException as exc: raise ValueError(f"Network error while polling video generation: {exc}") if response.status_code != 200: raise_video_api_error(response) try: payload = response.json() except Exception as exc: raise ValueError(f"Unable to parse video status response: {exc}") status = ( payload.get("status") or payload.get("state") or payload.get("data", {}).get("status") or payload.get("data", {}).get("state") or "" ).lower() if status in VIDEO_SUCCESS_STATES or (not status and extract_video_url(payload)): video_url = extract_video_url(payload) if not video_url: raise ValueError(f"Video generation completed but no video URL was returned: {payload}") return video_url if status in VIDEO_FAILURE_STATES: detail = payload.get("error") or payload.get("message") or payload raise ValueError(f"Video generation {generation_id} failed with status '{status}': {detail}") last_status = status or last_status time.sleep(max(interval, 1.0)) raise ValueError( f"Timed out after {timeout} seconds while waiting for video generation {generation_id}. " f"Last known status: {last_status}" ) def run_video_job(api_key: str, payload: dict, timeout: int = DEFAULT_VIDEO_TIMEOUT, interval: float = DEFAULT_VIDEO_POLL_INTERVAL): generation_id = submit_video_generation(api_key, payload) video_url = poll_video_generation(api_key, generation_id, timeout=timeout, interval=interval) return video_url, generation_id __all__ = [ "DEFAULT_VIDEO_POLL_INTERVAL", "DEFAULT_VIDEO_TIMEOUT", "VIDEO_GENERATIONS_ENDPOINT", "aspect_ratio_to_size", "build_auth_headers", "prepare_image_value", "run_video_job", "tensor_to_data_uri", "tensors_to_data_uri_list", "validate_api_key", ]