256 lines
8.6 KiB
Python
256 lines
8.6 KiB
Python
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",
|
|
]
|