#!/usr/bin/env python3
"""
Generate video with PolarQuant PQ5 LTX-2.3.

Usage:
    python generate_ltx.py --prompt "A cat playing piano" --output cat.mp4
    python generate_ltx.py --prompt "Ocean waves" --image ref.jpg --output ocean.mp4
"""
import argparse, os, subprocess, sys

def main():
    parser = argparse.ArgumentParser(description="LTX-2.3 Video Generation (PQ5)")
    parser.add_argument("--prompt", type=str, required=True)
    parser.add_argument("--image", type=str, help="Reference image for image-to-video")
    parser.add_argument("--output", type=str, default="output.mp4")
    parser.add_argument("--model-dir", type=str, default="./LTX-PQ5")
    parser.add_argument("--steps", type=int, default=50)
    parser.add_argument("--height", type=int, default=480)
    parser.add_argument("--width", type=int, default=704)
    parser.add_argument("--frames", type=int, default=97)
    args = parser.parse_args()

    model_dir = os.path.abspath(args.model_dir)
    code_dir = os.path.join(model_dir, "ltx_code")
    model_path = os.path.join(model_dir, "ltx-2.3-22b-dev.safetensors")

    if not os.path.exists(model_path):
        print("Error: Run `python setup.py` first.")
        sys.exit(1)

    print("=" * 60)
    print(f"  LTX-2.3 Video Generation (PQ5)")
    print(f"  Prompt: {args.prompt}")
    print(f"  Resolution: {args.width}x{args.height}, {args.frames} frames")
    print("=" * 60)

    # Use ltx-pipelines for inference
    cmd = [
        sys.executable, "-m", "ltx_pipelines.generate",
        "--model_path", model_path,
        "--prompt", args.prompt,
        "--output_path", args.output,
        "--height", str(args.height),
        "--width", str(args.width),
        "--num_frames", str(args.frames),
        "--num_inference_steps", str(args.steps),
    ]

    if args.image:
        cmd.extend(["--image_path", args.image])

    result = subprocess.run(cmd, cwd=code_dir)
    if os.path.exists(args.output):
        print(f"\nDone! {args.output} ({os.path.getsize(args.output)/1e6:.1f} MB)")
    else:
        print(f"\nFailed. Check LTX-2 docs: {code_dir}/packages/ltx-pipelines/README.md")

if __name__ == "__main__":
    main()