Image to video pipeline

import os
import sys
import time
from decimal import Decimal
from pathlib import Path

import httpx

API = os.environ.get("WINDPAINT_API_URL", "https://api.windpaint.ai").rstrip("/") + "/v1"
TERMINAL = {"completed", "failed", "nsfw", "canceled"}

api = httpx.Client(
    base_url=API,
    headers={"Authorization": f"Bearer {os.environ['WINDPAINT_API_KEY']}"},
    timeout=60,
)
out_dir = Path("out")
out_dir.mkdir(exist_ok=True)


def check(response: httpx.Response) -> dict:
    """Raise with the API's error code and message on any non-2xx response."""
    if response.is_success:
        return response.json()
    try:
        error = response.json()["error"]
    except (ValueError, KeyError):
        error = response.text
    if isinstance(error, dict):
        sys.exit(f"{response.status_code} {error['code']}: {error['message']} (request_id {error.get('request_id')})")
    sys.exit(f"{response.status_code}: {error}")


def submit(capability: str, body: dict) -> str:
    job = check(api.post(f"/generation/capabilities/{capability}", json=body))
    print(f"submitted {capability} {job['request_id']} ({job['credits_estimate']} credits held)")
    return job["request_id"]


def wait(request_id: str) -> dict:
    delay = 1.5
    while True:
        status = check(api.get(f"/generation/requests/{request_id}/status"))
        if status["status"] in TERMINAL:
            if status["status"] != "completed":
                sys.exit(f"{request_id} ended {status['status']}: {status['error']}")
            return status
        time.sleep(delay)
        delay = min(delay * 1.5, 15)


def download(asset_id: str, path: Path) -> Path:
    # A signed URL, valid 15 minutes. Fetch it without the Authorization header.
    url = check(api.get(f"/generation/assets/{asset_id}/url"))["data"]["url"]
    with httpx.stream("GET", url, timeout=300) as response:
        response.raise_for_status()
        with path.open("wb") as fh:
            for chunk in response.iter_bytes():
                fh.write(chunk)
    return path


# 1. Three candidate stills, submitted together, then waited on.
prompt = "a lighthouse on a rocky point at dusk, long exposure, waves, cinematic"
seeds = [11, 22, 33]
jobs = [
    submit("image.generate", {"prompt": prompt, "aspect_ratio": "16:9", "resolution": "1k", "seed": seed})
    for seed in seeds
]
stills = []
for seed, request_id in zip(seeds, jobs):
    asset = wait(request_id)["outputs"][0]
    path = download(asset["id"], out_dir / f"still-seed{seed}.png")
    stills.append(asset["id"])
    print(f"[{len(stills)}] {path}  ({asset['width']}x{asset['height']})")

# 2. Pick one.
choice = int(input(f"Animate which still? [1-{len(stills)}] ")) - 1
start_frame = stills[choice]

# 3. Estimate the clip and confirm before spending.
video = {
    "prompt": "slow push-in toward the lighthouse, waves breaking on the rocks, the beam starts to turn",
    "inputs": {"start_frame": [start_frame]},
    "aspect_ratio": "16:9",
    "resolution": "480p",
    "duration": 5,
}
estimate = check(api.post("/generation/estimate", json={"capability": "video.generate", **video}))["data"]
balance = check(api.get("/billing/balance"))["data"]
if estimate["credits"] is None:
    sys.exit(f"{estimate['model']} has no price for this tier")
print(f"clip: {estimate['model']} {estimate['width']}x{estimate['height']}, "
      f"{estimate['credits']} credits ({balance['available']} available)")
if Decimal(estimate["credits"]) > Decimal(balance["available"]):
    sys.exit("not enough credits; top up in Settings -> Billing")
if input("Submit? [y/N] ").lower() != "y":
    sys.exit("skipped")

# 4. Animate and download.
clip = wait(submit("video.generate", video))
path = download(clip["video"]["id"], out_dir / "clip.mp4")
print(f"saved {path}, charged {clip['credits']['actual']} credits")