# SPDX-License-Identifier: Apache-2.0 """Generate synchronized video/audio with Pika H3 Flex. Five task shapes share one model release: ``t2va`` text only ``i2va`` first image ``l2va`` last image ``fl2va`` first and last images ``ref2va`` ordered image/video/audio references The six-forward ladder is part of the adapter contract; this script deliberately does not expose a step-count flag. ``--lossless-accel`` turns on the existing training-free FastVideo switches (bit-identical kernels, reference memo, parallel VAE tiles) before the generator starts. ``--nvfp4-decoder`` and ``--nvfp4-text-encoder`` opt into the existing Blackwell FP4 weight paths; ``--weights nvfp4-qad`` (or ``nvfp4-qad-fp4-attention``) selects FastVideo's NVFP4 QAD release; ``--nvfp4-dit h3_dit`` loads any packed NVFP4 DiT export (e.g. FastVideo/Pika-H3-Flex-NVFP4-QAD). They change weights, not the six-forward sampler. Kernel/VSA work is separate. ``--nvfp4-attention`` turns on FlashAttention-4 NVFP4 attention (Q/K in NVFP4, P/V in BF16; Blackwell and FA4 only) with the named exact Q/K preconditioning, which must match the ``qk_transform`` the checkpoint was QAD-trained with (``smooth_k+hadamard`` for Pika-H3-Flex-NVFP4-QAD). ``--nvfp4-attention-bf16-blocks 45,49`` keeps those ``transformer_blocks`` on BF16 attention, as in training. ``none`` quantizes Q/K without preconditioning. Without ``--nvfp4-attention`` attention stays BF16. ``--resolution`` (480p/768p), ``--aspect`` (one of Pika's eleven buckets), and ``--duration`` (5-15 s) pick the bucket and the ``17n + 5`` frame count; ``--height``/``--width``/``--num-frames`` override them with explicit values. Launch with ``FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1`` on a FlashAttention-4 host (Hopper/Blackwell); without FA4 attention falls back to FA3/FA2 or, if neither is installed, to SDPA. ``--offload resident`` keeps every component on the GPU, which is how the GB200 timings in docs/inference/pika_h3_flex.md were measured; the default ``lazy`` mode loads components on demand for smaller devices and is slower. Examples: FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 \ python examples/inference/basic/basic_pika_h3_flex.py --task t2va --offload resident \ --prompt "a paper boat drifting down a rainy street" python examples/inference/basic/basic_pika_h3_flex.py --task t2va \ --resolution 480p --aspect 9:16 --duration 10 --prompt "a lighthouse at dusk" python examples/inference/basic/basic_pika_h3_flex.py --task fl2va \ --image first.png --last-image last.png --prompt "the flower opens" python examples/inference/basic/basic_pika_h3_flex.py --task ref2va \ --reference image:person.png --reference video:dance.mp4 \ --reference audio:voice.wav --prompt "the person dances to the music" """ from __future__ import annotations import argparse import os from pathlib import Path from typing import Any from fastvideo import VideoGenerator, envs from fastvideo.api import ( CompileConfig, ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig, ParallelismConfig, PipelineSelection, QuantizationConfig, SamplingConfig, ) from fastvideo.attention.utils.nvfp4_qk_precondition import QK_TRANSFORMS, parse_block_indices from fastvideo.models.vision_utils import load_image from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference from fastvideo.pipelines.basic.pika_h3_flex.buckets import ( PIKA_ASPECTS, PIKA_FPS, PIKA_MAX_DURATION, PIKA_MIN_DURATION, canvas_for_aspect, resolve_frames, ) # Training-free switches FastVideo already exposes. Workers inherit the # environment, so they must be set before VideoGenerator starts. LOSSLESS_ACCEL_ENV = { "FASTVIDEO_MINIMAX_H3_EXACT_KERNELS": "all", "FASTVIDEO_H3_REF2VA_MEMO_ENTRIES": "16", "FASTVIDEO_H3_VAE_TILE_PARALLEL": "1", "FASTVIDEO_ULYSSES_A2A": "auto", } TASKS = ("t2va", "i2va", "l2va", "fl2va", "ref2va") RESOLUTIONS = {"480p": "p480", "768p": "p768"} def _reference(value: str) -> MiniMaxH3Reference: try: media_type, source = value.split(":", 1) except ValueError as error: raise argparse.ArgumentTypeError("A reference must be image:PATH, video:PATH, or audio:PATH.") from error if media_type not in ("image", "video", "audio"): raise argparse.ArgumentTypeError(f"Unknown reference type {media_type!r}.") return MiniMaxH3Reference(source=source, media_type=media_type) WEIGHT_PRESETS = { "bf16": { "model_path": "Pika-Labs/Pika-H3-Flex" }, "nvfp4-qad": { "model_path": "FastVideo/Pika-H3-Flex-NVFP4-QAD", "nvfp4_dit": "h3_dit" }, "nvfp4-qad-fp4-attention": { "model_path": "FastVideo/Pika-H3-Flex-NVFP4-FP4Attn-QAD", "nvfp4_dit": "h3_dit", "nvfp4_attention": "smooth_k+hadamard", "nvfp4_attention_bf16_blocks": "45,49" }, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--weights", choices=tuple(WEIGHT_PRESETS), default="bf16", help="bf16: Pika's release; nvfp4-qad: FastVideo's NVFP4 QAD release; " "nvfp4-qad-fp4-attention: NVFP4 QAD with FP4 attention (both need Blackwell)") parser.add_argument("--model-path", help="local model directory or Hub repo (overrides --weights)") parser.add_argument("--task", choices=TASKS, default="t2va") parser.add_argument("--prompt", required=True) parser.add_argument("--image", help="first frame (i2va/fl2va)") parser.add_argument("--last-image", help="last frame (l2va/fl2va)") parser.add_argument("--reference", action="append", type=_reference, default=[], help="ordered reference (repeatable; ref2va)") parser.add_argument("--resolution", choices=tuple(RESOLUTIONS), default="768p", help="Pika resolution tier") parser.add_argument("--aspect", choices=PIKA_ASPECTS, default="16:9", help="bucket aspect ratio (W:H)") parser.add_argument("--duration", type=float, default=5.0, help=f"seconds ({PIKA_MIN_DURATION:g}-{PIKA_MAX_DURATION:g}); rounded up to the 17n+5 grid") parser.add_argument("--height", type=int, help="explicit canvas height (with --width; overrides the bucket)") parser.add_argument("--width", type=int, help="explicit canvas width (with --height; overrides the bucket)") parser.add_argument("--num-frames", type=int, help="explicit frame count (overrides --duration)") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--num-gpus", type=int, default=1) parser.add_argument("--compile", action="store_true", help="torch.compile the DiT") parser.add_argument("--lossless-accel", action="store_true", help="enable bit-identical speed-ups") parser.add_argument("--offload", choices=("lazy", "resident"), default="lazy", help="lazy: load components on demand and offload the encoders/VAE (fits smaller devices); " "resident: keep everything on the GPU (fastest, about 105-130 GiB on GB200)") parser.add_argument("--low-vram", action="store_true", help="stream the DiT layerwise (lazy offload only)") parser.add_argument("--nvfp4-decoder", help="Blackwell-only NVFP4 H3 decoder (quality-36L or fast-8L decoder-nvfp4.safetensors)") parser.add_argument("--nvfp4-text-encoder", help="Blackwell-only serialized NVFP4 Qwen3-VL text_encoder directory") parser.add_argument("--nvfp4-dit", choices=("h3_dit", "h3_dit_ffn"), help="load a packed NVFP4 DiT export with this layer profile (h3_dit: attention + FFN linears)") parser.add_argument("--nvfp4-attention", choices=tuple(QK_TRANSFORMS), help="FA4 NVFP4 Q/K attention with this exact Q/K preconditioning (default: BF16 attention)") parser.add_argument("--nvfp4-attention-bf16-blocks", help="comma-separated transformer_blocks indices that keep BF16 attention, e.g. 45,49") parser.add_argument("--output", default="outputs/pika_h3_flex") args = parser.parse_args() for name, value in WEIGHT_PRESETS[args.weights].items(): if getattr(args, name) is None: setattr(args, name, value) if args.nvfp4_attention_bf16_blocks and not args.nvfp4_attention: parser.error("--nvfp4-attention-bf16-blocks needs --nvfp4-attention") try: parse_block_indices(args.nvfp4_attention_bf16_blocks) except ValueError as error: parser.error(f"--nvfp4-attention-bf16-blocks: {error}") if args.task in ("i2va", "fl2va") and not args.image: parser.error(f"--task {args.task} needs --image") if args.task in ("l2va", "fl2va") and not args.last_image: parser.error(f"--task {args.task} needs --last-image") if args.task == "ref2va" and not args.reference: parser.error("--task ref2va needs at least one --reference") if args.task != "ref2va" and args.reference: parser.error("--reference is only valid with --task ref2va") if (args.height is None) != (args.width is None): parser.error("--height and --width must be passed together") if args.low_vram and args.offload == "resident": parser.error("--low-vram streams the DiT and cannot be combined with --offload resident") return args def _offload(args: argparse.Namespace) -> OffloadConfig: if args.offload == "resident": return OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False, vae=False, pin_cpu_memory=False, lazy_module_load=False) return OffloadConfig( dit=False, dit_layerwise=args.low_vram, text_encoder=True, vae=True, pin_cpu_memory=not args.low_vram, lazy_module_load=True, ) def _canvas(args: argparse.Namespace) -> tuple[int, int]: if args.height is not None: return args.height, args.width return canvas_for_aspect(args.aspect, RESOLUTIONS[args.resolution]) def _num_frames(args: argparse.Namespace) -> int: if args.num_frames is not None: return args.num_frames return resolve_frames(round(args.duration * PIKA_FPS)) def _inputs(args: argparse.Namespace) -> InputConfig: if args.task == "ref2va": return InputConfig(references=args.reference) first = args.image if args.task in ("i2va", "fl2va") else None last = args.last_image if args.task in ("l2va", "fl2va") else None # The stage loads ``image_path`` itself but requires ``last_image`` as a PIL image. return InputConfig(image_path=first, last_image=load_image(last) if last else None) def _pipeline(args: argparse.Namespace) -> PipelineSelection: components = ComponentConfig() if args.task == "ref2va": components.override_pipeline_cls_name = "PikaH3FlexRef2VAModularPipeline" if args.nvfp4_decoder: components.h3_vae_nvfp4_decoder = args.nvfp4_decoder if args.nvfp4_text_encoder: components.text_encoder_weights = args.nvfp4_text_encoder return PipelineSelection( workload_type="t2v" if args.task == "t2va" else "i2v", components=components, vae_tiling=True, ) def _apply_nvfp4_attention(args: argparse.Namespace) -> None: """Workers inherit the environment, so set it before VideoGenerator starts.""" if not args.nvfp4_attention: return envs.FASTVIDEO_NVFP4_FA4.set(True) envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.set(args.nvfp4_attention) envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.set(args.nvfp4_attention_bf16_blocks or "") envs.setdefault_external("CUTE_DSL_ENABLE_TVM_FFI", "1") def main() -> None: args = parse_args() _apply_nvfp4_attention(args) if args.lossless_accel: for key, value in LOSSLESS_ACCEL_ENV.items(): os.environ.setdefault(key, value) height, width = _canvas(args) num_frames = _num_frames(args) print(f"Pika H3 Flex {args.task}: {width}x{height}, {num_frames} frames ({num_frames / PIKA_FPS:.2f} s)") output_path = Path(args.output) / f"pika_h3_flex_{args.task}.mp4" output_path.parent.mkdir(parents=True, exist_ok=True) offload = _offload(args) generator = VideoGenerator.from_config( GeneratorConfig( model_path=args.model_path, engine=EngineConfig( num_gpus=args.num_gpus, use_fsdp_inference=args.num_gpus > 1, parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), offload=offload, compile=CompileConfig(enabled=args.compile), quantization=(QuantizationConfig(transformer_quant="NVFP4", layer_profile=args.nvfp4_dit) if args.nvfp4_dit else None), ), pipeline=_pipeline(args), )) try: result = generator.generate( GenerationRequest( prompt=args.prompt, negative_prompt="", inputs=_inputs(args), sampling=SamplingConfig( height=height, width=width, num_frames=num_frames, fps=24, num_inference_steps=6, guidance_scale=1.0, batch_cfg=False, seed=args.seed, ), output=OutputConfig( output_path=str(output_path), save_video=True, return_frames=False, ), )) print(f"Output written to: {result.video_path}") finally: generator.shutdown() if __name__ == "__main__": main()