scripts/run-ltx-cloud.py

"""
run-ltx-cloud.py — Run LTX-2.3 video inference on HF Jobs cloud GPU
Usage:
  hf jobs run --flavor t4-medium --image python:3.12 \
    "python run-ltx-cloud.py --scenes scenes.json --output final.mp4"
"""
import json, os, sys, subprocess, time, shutil, glob, argparse, gc

parser = argparse.ArgumentParser()
parser.add_argument("--scenes", default="scenes.json")
parser.add_argument("--output", default="final.mp4")
parser.add_argument("--bucket", default="hf://buckets/kritnatee/Creator360.Studio-storage")
parser.add_argument("--steps", type=int, default=20)
parser.add_argument("--frames", type=int, default=49)
args = parser.parse_args()

# Step 1: Install system + Python deps
print("[1/5] Installing deps...")
# Check if ffmpeg is available, install if not
if subprocess.run(["which", "ffmpeg"], capture_output=True).returncode != 0:
    print("  Installing ffmpeg...")
    try:
        subprocess.run(["apt-get", "update", "-qq"], check=True, capture_output=True)
        subprocess.run(["apt-get", "install", "-y", "-qq", "ffmpeg"], check=True, capture_output=True)
    except:
        print("  WARN: could not install ffmpeg; video assembly may fail")

subprocess.run(["pip", "install", "-q",
    "diffusers", "transformers", "accelerate",
    "safetensors", "imageio[ffmpeg]", "psutil", "huggingface_hub[hf_xet]"], check=True)

# Step 2: Download assets
print("[2/5] Downloading from bucket...")
os.makedirs("images", exist_ok=True)
subprocess.run(["hf", "sync", f"{args.bucket}/images", "images"], check=True)
if not os.path.exists(args.scenes):
    subprocess.run(["hf", "sync", f"{args.bucket}/scripts/{args.scenes}", "."], capture_output=True)

image_files = sorted(sum([glob.glob(f"images/**/*{e}", recursive=True)
    for e in ("png", "jpg", "jpeg", "webp")], []))

# Step 3: Load plan
print("[3/5] Loading scene plan...")
if os.path.exists(args.scenes):
    with open(args.scenes) as f:
        scenes = json.load(f)["scenes"]
else:
    scenes = [{"scene": i+1, "start_image": img, "transition_in": "fade",
               "transition_out": "fade" if i == len(image_files)-1 else "crossfade",
               "duration": 3.0, "mood": "calm"}
              for i, img in enumerate(image_files)]
print(f"  {len(scenes)} scenes")

# Step 4: Run LTX-2.3 for each scene
print("[4/5] LTX-2.3 inference...")
gpu_info = os.popen("nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null").read().strip()
print(f"  GPU: {gpu_info}")

from diffusers import LTXPipeline
import torch
from PIL import Image
import numpy as np
import imageio

os.makedirs("clips", exist_ok=True)

# Load model with appropriate memory settings
print("  Loading LTX-Video (22B)...")
import psutil
ram_gb = psutil.virtual_memory().total / (1024**3)
vram_gb = 0
try:
    vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
except:
    pass
print(f"  RAM: {ram_gb:.0f} GB | VRAM: {vram_gb:.0f} GB")

# Try loading; fallback to 8-bit if RAM is tight
load_kwargs = {"torch_dtype": torch.bfloat16}
if ram_gb < 35:
    print("  Low RAM: enabling 8-bit + CPU offload")
    load_kwargs["torch_dtype"] = torch.float16

pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", **load_kwargs)
pipe.enable_model_cpu_offload()
pipe.vae.enable_slicing()

for i, scene in enumerate(scenes):
    img_ref = scene.get("start_image", "")
    img_path = None

    # Find image
    if img_ref and os.path.exists(img_ref):
        img_path = img_ref
    elif img_ref:
        matches = [f for f in image_files if img_ref in f]
        if matches:
            img_path = matches[0]
    if not img_path and i < len(image_files):
        img_path = image_files[i]

    if not img_path:
        print(f"  Scene {scene['scene']}: SKIP")
        continue

    print(f"  Scene {scene['scene']}: {os.path.basename(img_path)}   ", end="")
    sys.stdout.flush()

    t0 = time.time()
    try:
        input_img = Image.open(img_path).convert("RGB")
        prompt = scene.get("narrative", f"Cinematic {scene.get('mood','calm')} scene")

        frames = pipe(
            image=input_img,
            prompt=prompt,
            negative_prompt="blurry, low quality",
            num_frames=args.frames,
            width=704, height=480,
            num_inference_steps=args.steps,
            generator=torch.Generator(device="cpu").manual_seed(42 + i),
        ).frames[0]

        clip_path = f"clips/scene_{i:03d}.mp4"
        imageio.mimwrite(clip_path, frames, fps=24, codec="libx264", quality=8)
        print(f" OK ({time.time()-t0:.0f}s)")
    except Exception as e:
        print(f" FAIL: {e}")
        # Fallback: static image
        try:
            img_np = np.array(Image.open(img_path).resize((704, 480)))
            fallback = np.tile(img_np[None], (args.frames, 1, 1, 1))
            imageio.mimwrite(f"clips/scene_{i:03d}.mp4", fallback, fps=24, codec="libx264")
            print("  Used static fallback")
        except:
            pass

torch.cuda.empty_cache()
gc.collect()

# Step 5: Assemble & upload
print("[5/5] Assembling video + upload...")
clips = sorted(glob.glob("clips/scene_*.mp4"))
if not clips:
    print("ERROR: no clips"); sys.exit(1)

concat = "\n".join(f"file '{os.path.abspath(c)}'" for c in clips)
with open("_concat.txt", "w") as f:
    f.write(concat)

subprocess.run(["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", "_concat.txt",
    "-c:v", "libx264", "-preset", "slow", "-crf", "18",
    "-pix_fmt", "yuv420p", "-movflags", "+faststart",
    args.output], check=True)

subprocess.run(["hf", "sync", os.path.abspath(args.output),
    f"{args.bucket}/video/{args.output}"], check=True)

size_mb = os.path.getsize(args.output) / 1024 / 1024
print(f"\nDONE: {args.output} ({size_mb:.1f} MB)")
print(f"Bucket: {args.bucket}/video/{args.output}")