# SPDX-License-Identifier: Apache-2.0 # Copyright 2026 Mellis, Inc. """Shared loading, memory management and output checks for the two examples.""" import argparse import hashlib import json import os from pathlib import Path import re import subprocess import time os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True") CODE_REVISION = "d0b9b1b5a540e5f9df2f8ec236347bf9c5aa7d1f" BUCKETS = {"480p": (832, 480), "768p": (1344, 768)} def arguments(task): parser = argparse.ArgumentParser(description=f"Pika H3 Flex {task} example (16:9)") parser.add_argument("--repo", default="Pika-Labs/Pika-H3-Flex") parser.add_argument("--revision", required=True, help="Immutable 40-character HF commit SHA") parser.add_argument("--snapshot", type=Path, help="Optional complete local snapshot of that revision") parser.add_argument("--assets-dir", type=Path, default=Path(__file__).parent / "assets", help="Directory containing fl2va/ and ref2va/ assets; missing assets download from the same HF revision") parser.add_argument("--resolution", choices=BUCKETS, default="480p") parser.add_argument("--duration", type=int, choices=range(5, 16), default=15 if task == "fl2va" else 10) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--output", type=Path, default=Path("outputs") / f"{task}.mp4") parser.add_argument("--resident-transformer", action="store_true", help="Disable transformer block offload; uses more GPU memory") parser.add_argument("--text-encoder-offload", action="store_true", help="Stream text/vision encoder layers from CPU for larger reference inputs") args = parser.parse_args() if not re.fullmatch(r"[a-f0-9]{40}", args.revision): parser.error("--revision must be an immutable 40-character HF commit SHA") return args def asset(args, task, filename): path = args.assets_dir / task / filename if path.is_file(): return path from huggingface_hub import hf_hub_download return Path(hf_hub_download(args.repo, f"examples/assets/{task}/{filename}", revision=args.revision)) def offload_other_components(hooks, model_id, model, execution_device): # A full text encoder and transformer cannot coexist on an H100 80GB. # Move all other resident components aside when the next component starts. # The manager itself excludes group-offloaded components from this list. return hooks def build_pipeline(args, task): import torch from diffusers import ComponentsManager, ModularPipeline if not torch.cuda.is_available(): raise RuntimeError("An NVIDIA CUDA GPU is required for these examples") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = True manager = ComponentsManager() manager.enable_auto_cpu_offload(device="cuda:0", offload_strategy=offload_other_components) source = str(args.snapshot) if args.snapshot else args.repo pipe = ModularPipeline.from_pretrained(source, revision=args.revision, workflow=task, components_manager=manager) pipe.load_components( pretrained_model_name_or_path=source, revision=args.revision, dtype={"vae": torch.float32, "audio_vae": torch.float32, "default": torch.bfloat16}, ) pipe.load_lora_weights(source, revision=args.revision, weight_name=f"{task}/adapter.safetensors") if args.text_encoder_offload: from accelerate import cpu_offload from accelerate.hooks import remove_hook_from_module # The pipeline calls text_encoder.model directly and does not use the # language-model head. Let Accelerate place those actual computation # layers, rather than moving the entire encoder onto the GPU at once. remove_hook_from_module(pipe.text_encoder) cpu_offload(pipe.text_encoder.model, execution_device=torch.device("cuda:0")) transformer = pipe.transformer_ref if task == "ref2va" else pipe.transformer def release_other_components(module, inputs): # Run before group-offload's wrapped forward: once its first group is on # CUDA, the manager can otherwise mistake the whole model for resident. for hook in manager.model_hooks: if hook.model is not module and hook.model.device.type == "cuda": hook.offload() transformer.register_forward_pre_hook(release_other_components, prepend=True) if not args.resident_transformer: transformer.enable_group_offload( onload_device=torch.device("cuda:0"), offload_device=torch.device("cpu"), offload_type="block_level", num_blocks_per_group=1, use_stream=True, record_stream=False, ) return pipe def generate(args, task, media_kwargs, inputs): import torch import diffusers from diffusers.utils.export_utils import encode_video prompt_path = asset(args, task, "prompt.txt") prompt = prompt_path.read_text(encoding="utf-8").strip() width, height = BUCKETS[args.resolution] expected_frames = args.duration * 24 while expected_frames % 17 != 5: expected_frames += 1 start = time.monotonic() pipe = build_pipeline(args, task) loaded = time.monotonic() torch.cuda.reset_peak_memory_stats() result = pipe(prompt=prompt, height=height, width=width, num_frames=args.duration * 24, generator=torch.Generator("cuda:0").manual_seed(args.seed), output=["videos", "audio", "sampling_rate"], **media_kwargs) torch.cuda.synchronize() generated = time.monotonic() frames = result["videos"][0] audio = result["audio"][0].detach().float().cpu() assert len(frames) == expected_frames, (len(frames), expected_frames) assert frames[0].size == (width, height), frames[0].size assert torch.isfinite(audio).all(), "Output audio contains NaN or Inf" assert result["sampling_rate"] == 32000 and audio.shape[0] == 2 args.output.parent.mkdir(parents=True, exist_ok=True) encode_video(frames, fps=24, output_path=str(args.output), audio=audio, audio_sample_rate=result["sampling_rate"]) probe = json.loads(subprocess.check_output( ["ffprobe", "-v", "error", "-show_streams", "-show_format", "-of", "json", str(args.output)], text=True)) video = next(s for s in probe["streams"] if s["codec_type"] == "video") sound = next(s for s in probe["streams"] if s["codec_type"] == "audio") assert (video["width"], video["height"]) == (width, height) assert int(video["nb_frames"]) == expected_frames and video["r_frame_rate"] == "24/1" assert int(sound["sample_rate"]) == 32000 and sound["channels"] == 2 record = { "repo": args.repo, "revision": args.revision, "diffusers_code_revision": CODE_REVISION, "diffusers_version": diffusers.__version__, "task": task, "seed": args.seed, "resolution": args.resolution, "width": width, "height": height, "requested_duration": args.duration, "frames": expected_frames, "fps": 24, "transformer_offload": not args.resident_transformer, "text_encoder_offload": args.text_encoder_offload, "load_seconds": loaded - start, "generation_seconds": generated - loaded, "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30, "input_sha256": {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in [prompt_path, *inputs]}, "output_sha256": hashlib.sha256(args.output.read_bytes()).hexdigest(), "media": probe, } args.output.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n") print(json.dumps({"output": str(args.output), "frames": expected_frames, "generation_seconds": record["generation_seconds"], "peak_allocated_gib": record["peak_allocated_gib"]}), flush=True)