ResembleAI/Dramabox

1#!/usr/bin/env python32"""Warm A2Vid driver — builds :class:`A2VidPipelineTwoStage` once and runs3multiple scenes back-to-back without reloading Gemma, the 22B DiT, VAEs or4the upsampler between calls. Saves ~85 s of model-load time per scene.5 6Scenes are declared in a JSON manifest:7 8    {9      "scenes": [10        {11          "name": "scene01",12          "audio": "scene01.wav",13          "prompt": "...",14          "num_frames": 241,15          "tail_from": null                    # no conditioning for first scene16        },17        {18          "name": "scene02",19          "audio": "scene02.wav",20          "prompt": "...",21          "num_frames": 201,22          "tail_from": "scene01",              # pin first 24 frames to scene01 tail23          "tail_seconds": 1.0,24          "tail_strength": 0.725        }26      ]27    }28 29Each scene writes <out>/<name>.mp4 and its tail-frames get extracted if a30later scene references it.31"""32import argparse33import json34import logging35import os36import subprocess37import sys38import time39from pathlib import Path40 41import torch42 43 44def extract_tail_frames(mp4: Path, out_prefix: Path, seconds: float, fps: float) -> list[Path]:45    """Extract the last ``seconds`` seconds of ``mp4`` as PNGs starting at index 0."""46    dur = float(subprocess.check_output(47        ["ffprobe", "-v", "error", "-select_streams", "v:0",48         "-show_entries", "stream=duration", "-of", "csv=p=0", str(mp4)],49    ).decode().strip())50    start = max(0.0, dur - seconds - 0.05)51    n_frames = int(round(seconds * fps))52    # Clean stale53    for p in out_prefix.parent.glob(f"{out_prefix.name}_*.png"):54        p.unlink()55    subprocess.run(56        ["ffmpeg", "-y", "-ss", f"{start:.3f}", "-i", str(mp4),57         "-vf", f"fps={fps}", "-frames:v", str(n_frames),58         "-start_number", "0", f"{out_prefix}_%03d.png",59         "-loglevel", "error"],60        check=True,61    )62    return sorted(out_prefix.parent.glob(f"{out_prefix.name}_*.png"))63 64 65def main():66    ap = argparse.ArgumentParser()67    ap.add_argument("--manifest", required=True, help="JSON scene manifest")68    ap.add_argument("--out-dir", required=True)69    ap.add_argument("--checkpoint-path", required=True)70    ap.add_argument("--gemma-root", required=True)71    ap.add_argument("--spatial-upsampler-path", required=True)72    ap.add_argument("--distilled-lora", required=True)73    ap.add_argument("--quantization", default="fp8-cast",74                    choices=["fp8-cast", "none"])75    ap.add_argument("--bnb-4bit", action="store_true", default=True,76                    help="Load Gemma via bnb-4bit path (default on).")77    ap.add_argument("--no-bnb-4bit", dest="bnb_4bit", action="store_false")78    ap.add_argument("--seed", type=int, default=42)79    ap.add_argument("--num-inference-steps", type=int, default=30)80    ap.add_argument("--height", type=int, default=512)81    ap.add_argument("--width", type=int, default=768)82    ap.add_argument("--frame-rate", type=float, default=24.0)83    # Defaults for guider params (match a2vid_two_stage CLI defaults)84    ap.add_argument("--cfg-scale", type=float, default=2.5)85    ap.add_argument("--stg-scale", type=float, default=1.0)86    ap.add_argument("--rescale-scale", type=float, default=0.7)87    ap.add_argument("--modality-scale", type=float, default=2.5)88    ap.add_argument("--negative-prompt", default=89                    "low quality, worst quality, blurry, distorted, artifacts, watermark, text, caption")90    args = ap.parse_args()91 92    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")93    out_dir = Path(args.out_dir)94    out_dir.mkdir(parents=True, exist_ok=True)95 96    # Import after argparse so --help is instant.97    from ltx_core.loader.registry import Registry98    from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number99    from ltx_core.quantization import QuantizationPolicy100    from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage101    from ltx_pipelines.utils.args import ImageConditioningInput102    from ltx_pipelines.utils.blocks import PromptEncoder103    from ltx_core.components.guiders import MultiModalGuiderParams104    from ltx_pipelines.utils.media_io import encode_video105 106    quant = QuantizationPolicy.fp8_cast() if args.quantization == "fp8-cast" else None107    registry = Registry()108 109    logging.info("Building warm A2Vid pipeline (loads Gemma, VAEs, DiT, upsampler)...")110    t0 = time.time()111    pipeline = A2VidPipelineTwoStage(112        checkpoint_path=args.checkpoint_path,113        distilled_lora=[(args.distilled_lora, 1.0, None)] if False else [],  # see __init__114        spatial_upsampler_path=args.spatial_upsampler_path,115        gemma_root=args.gemma_root,116        loras=(),117        quantization=quant,118        registry=registry,119    )120    # Replace the pipeline's PromptEncoder with a warm+bnb one so subsequent121    # calls skip the Gemma load. A2VidPipelineTwoStage stored it as122    # self.prompt_encoder.123    logging.info("Replacing PromptEncoder with warm + bnb-4bit variant...")124    pipeline.prompt_encoder = PromptEncoder(125        checkpoint_path=args.checkpoint_path,126        gemma_root=args.gemma_root,127        dtype=torch.bfloat16,128        device=pipeline.device,129        registry=registry,130        warm=True,131        use_bnb_4bit=args.bnb_4bit,132    )133    logging.info(f"Pipeline ready in {time.time() - t0:.1f}s")134 135    manifest = json.loads(Path(args.manifest).read_text())136    tiling = TilingConfig.default()137    mp4_paths: dict[str, Path] = {}138 139    for scene in manifest["scenes"]:140        name = scene["name"]141        mp4 = out_dir / f"{name}.mp4"142        mp4_paths[name] = mp4143        if mp4.exists():144            logging.info(f"[{name}] skipping — already exists")145            continue146 147        num_frames = int(scene["num_frames"])148        # Build image conditioning from an earlier scene's tail, if specified.149        images: list[ImageConditioningInput] = []150        tail_from = scene.get("tail_from")151        if tail_from:152            src_mp4 = mp4_paths.get(tail_from)153            if src_mp4 is None or not src_mp4.exists():154                raise RuntimeError(f"scene {name} needs tail from {tail_from} which hasn't been generated")155            secs = float(scene.get("tail_seconds", 1.0))156            strength = float(scene.get("tail_strength", 0.7))157            prefix = out_dir / f"{tail_from}_tail"158            logging.info(f"[{name}] extracting tail ({secs}s @ {args.frame_rate}fps) from {src_mp4.name}")159            tail_pngs = extract_tail_frames(src_mp4, prefix, secs, args.frame_rate)160            for i, png in enumerate(tail_pngs):161                images.append(ImageConditioningInput(str(png), i, strength))162 163        logging.info(f"[{name}] generating {num_frames} frames, {len(images)} conditioning images")164        t1 = time.time()165        video, audio = pipeline(166            prompt=scene["prompt"],167            negative_prompt=args.negative_prompt,168            seed=args.seed,169            height=args.height,170            width=args.width,171            num_frames=num_frames,172            frame_rate=args.frame_rate,173            num_inference_steps=args.num_inference_steps,174            video_guider_params=MultiModalGuiderParams(175                cfg_scale=args.cfg_scale,176                stg_scale=args.stg_scale,177                rescale_scale=args.rescale_scale,178                modality_scale=args.modality_scale,179            ),180            images=images,181            tiling_config=tiling,182            audio_path=scene["audio"],183            audio_start_time=0.0,184            audio_max_duration=num_frames / args.frame_rate,185        )186        encode_video(187            video=video, fps=args.frame_rate, audio=audio,188            output_path=str(mp4),189            video_chunks_number=get_video_chunks_number(num_frames, tiling),190        )191        logging.info(f"[{name}] done in {time.time() - t1:.1f}s -> {mp4}")192 193    logging.info("All scenes done.")194 195 196if __name__ == "__main__":197    sys.exit(main())198