Happy Horse 1.0

import json
import os
import time

import requests

MODEL = "happyhorse-1.0"
PROMPT = (
    "A miniature city built from cardboard and bottle caps comes alive at night. "
    "A tiny cardboard train rolls through the streets while warm lights glow from "
    "paper windows, physically realistic motion, detailed tabletop cinematography."
)
SECONDS = "3"
SIZE = "1280x720"
BASE_URL = "https://api.cometapi.com"
OUTPUT_DIR = "./output"
POLL_INTERVAL_SECONDS = 10
RETRY_DELAY_SECONDS = 5
MAX_CREATE_ATTEMPTS = 5
TERMINAL_STATUSES = {"completed", "failed", "error", "success"}
SUCCESS_STATUSES = {"completed", "success"}

try:
    COMETAPI_KEY = os.environ["COMETAPI_KEY"]
except KeyError:
    print("Set COMETAPI_KEY before running this example.")
    raise SystemExit(0)

headers = {"Authorization": f"Bearer {COMETAPI_KEY}"}

def is_progress_complete(progress):
    if isinstance(progress, (int, float)):
        return progress >= 100
    if isinstance(progress, str):
        try:
            return float(progress.rstrip("%")) >= 100
        except ValueError:
            return False
    return False

def is_transient_status(status_code):
    return status_code == 429 or 500 <= status_code < 600

def create_task():
    fields = {
        "model": (None, MODEL),
        "prompt": (None, PROMPT),
        "seconds": (None, SECONDS),
        "size": (None, SIZE),
    }
    for attempt in range(1, MAX_CREATE_ATTEMPTS + 1):
        try:
            response = requests.post(
                f"{BASE_URL}/v1/videos",
                headers=headers,
                files=fields,
                timeout=30,
            )
        except requests.RequestException as error:
            if attempt == MAX_CREATE_ATTEMPTS:
                raise
            print(f"Create request failed: {error}, retrying...")
            time.sleep(RETRY_DELAY_SECONDS)
            continue

        if response.ok:
            return response
        if not is_transient_status(response.status_code) or attempt == MAX_CREATE_ATTEMPTS:
            response.raise_for_status()
        print(f"Create request returned {response.status_code}, retrying...")
        time.sleep(RETRY_DELAY_SECONDS)
    raise SystemExit("Failed to create task.")

def get_task(task_id):
    while True:
        try:
            response = requests.get(
                f"{BASE_URL}/v1/videos/{task_id}",
                headers=headers,
                timeout=15,
            )
        except requests.RequestException as error:
            print(f"Status request failed: {error}, retrying...")
            time.sleep(RETRY_DELAY_SECONDS)
            continue

        if response.ok:
            return response
        if not is_transient_status(response.status_code):
            response.raise_for_status()
        print(f"Status request returned {response.status_code}, retrying...")
        time.sleep(RETRY_DELAY_SECONDS)

def download_video(content_url, output_path):
    while True:
        try:
            with requests.get(content_url, headers=headers, timeout=120, stream=True) as response:
                if response.ok:
                    with open(output_path, "wb") as output_file:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                output_file.write(chunk)
                    return
                if not is_transient_status(response.status_code):
                    response.raise_for_status()
                print(f"Content request returned {response.status_code}, retrying...")
        except requests.RequestException as error:
            print(f"Content request failed: {error}, retrying...")
        time.sleep(RETRY_DELAY_SECONDS)

create_response = create_task()
create_result = create_response.json()

task_id = create_result.get("id") or create_result.get("task_id")
if not task_id:
    print(json.dumps(create_result, indent=2))
    raise SystemExit("No task id returned.")

print(f"Task created: {task_id}")
print(f"Initial status: {create_result.get('status')}")

while True:
    task_response = get_task(task_id)
    task = task_response.json()
    status = str(task.get("status") or "unknown")
    normalized_status = status.lower()
    progress = task.get("progress")
    should_download = normalized_status in SUCCESS_STATUSES or (
        normalized_status == "unknown" and is_progress_complete(progress)
    )

    print(f"Status: {status}, progress: {progress}")

    if should_download or normalized_status in TERMINAL_STATUSES:
        if not should_download:
            print(json.dumps(task, indent=2))
            raise SystemExit(1)

        content_url = f"{BASE_URL}/v1/videos/{task_id}/content"
        output_path = os.path.join(OUTPUT_DIR, f"{task_id}.mp4")
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        download_video(content_url, output_path)

        print(f"Video URL: {task.get('video_url') or ''}")
        print(f"Content endpoint: {content_url}")
        print(f"Saved to {output_path}")
        print(f"File size: {os.path.getsize(output_path)} bytes")
        break

    time.sleep(POLL_INTERVAL_SECONDS)