batch_generation.py

"""
Seedance 2.0 – Batch Video Generation
======================================
Generate multiple videos in parallel using Python's ThreadPoolExecutor.
Useful for bulk content creation, A/B testing prompts, or dataset generation.

Get your API key: https://www.atlascloud.ai/seedance-2?utm_source=github&utm_campaign=free-seedance-2-api
"""

import os
import time
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass

API_KEY = os.environ.get("ATLASCLOUD_API_KEY", "your_api_key_here")
BASE_URL = "https://api.atlascloud.ai"


@dataclass
class VideoJob:
    prompt: str
    model: str = "bytedance/seedance-2.0/text-to-video-fast"
    width: int = 1280
    height: int = 720
    duration: int = 5
    output_name: str = "output"


def submit_job(job: VideoJob) -> tuple[str, str]:
    """Submit a single generation job. Returns (job_name, prediction_id)."""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
    }
    response = requests.post(
        f"{BASE_URL}/api/v1/model/generateVideo",
        headers=headers,
        json={
            "model": job.model,
            "input": {
                "prompt": job.prompt,
                "width": job.width,
                "height": job.height,
                "duration": job.duration,
            },
        },
        timeout=30,
    )
    response.raise_for_status()
    prediction_id = response.json()["data"]["id"]
    return job.output_name, prediction_id


def poll_until_done(output_name: str, prediction_id: str) -> tuple[str, str]:
    """Poll until the job completes. Returns (output_name, video_url)."""
    poll_url = f"{BASE_URL}/api/v1/model/prediction/{prediction_id}"
    poll_headers = {"Authorization": f"Bearer {API_KEY}"}
    interval = 3

    while True:
        result = requests.get(poll_url, headers=poll_headers, timeout=15).json()["data"]
        if result["status"] in ("completed", "succeeded"):
            return output_name, result["outputs"][0]
        if result["status"] == "failed":
            raise RuntimeError(f"[{output_name}] Failed: {result.get('error')}")
        time.sleep(interval)
        interval = min(interval * 1.3, 10)


def batch_generate(jobs: list[VideoJob], max_workers: int = 5) -> dict[str, str]:
    """
    Generate multiple videos in parallel.

    Args:
        jobs: List of VideoJob objects
        max_workers: Maximum concurrent generations (respect rate limits)

    Returns:
        Dict mapping output_name → video_url
    """
    results = {}

    # Submit all jobs
    print(f"Submitting {len(jobs)} jobs...")
    pending = []
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(submit_job, job): job for job in jobs}
        for future in as_completed(futures):
            try:
                output_name, prediction_id = future.result()
                pending.append((output_name, prediction_id))
                print(f"  ✓ Submitted: {output_name} → {prediction_id}")
            except Exception as e:
                print(f"  ✗ Submit error: {e}")

    # Poll all jobs in parallel
    print(f"\nPolling {len(pending)} jobs for completion...")
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(poll_until_done, name, pid): name
            for name, pid in pending
        }
        for future in as_completed(futures):
            try:
                output_name, video_url = future.result()
                results[output_name] = video_url
                print(f"  ✓ Done: {output_name} → {video_url}")
            except Exception as e:
                print(f"  ✗ Error: {e}")

    return results


if __name__ == "__main__":
    # Example: Generate 5 product videos in parallel
    products = [
        ("red_sneaker", "A red high-top sneaker rotating slowly on white background, studio lighting"),
        ("blue_bag", "A navy blue leather handbag on marble surface, luxury product photography"),
        ("watch", "A stainless steel watch face with moving hands, macro lens, dramatic lighting"),
        ("sunglasses", "Sunglasses with mirrored lenses on sandy beach, golden hour, commercial"),
        ("headphones", "Black wireless headphones floating on clean background, tech product shot"),
    ]

    jobs = [
        VideoJob(
            prompt=prompt,
            output_name=name,
            model="bytedance/seedance-2.0/text-to-video-fast",
            width=1080,
            height=1080,
            duration=5,
        )
        for name, prompt in products
    ]

    results = batch_generate(jobs, max_workers=5)

    print("\n=== Results ===")
    for name, url in results.items():
        print(f"{name}: {url}")

    # Estimated cost: 5 videos × 5 sec × $0.081 = $2.025
    print(f"\nEstimated cost: ${5 * 5 * 0.081:.3f}")