[feat]: add Pika H3 Flex inference support
Add Pika H3 Flex (MiniMax-H3 plus the released fl2va/ref2va rank-128 adapters) as a FastVideo model family: registry/config/presets, bucket geometry with tiered reference sizing, automatic adapter selection and manifest validation, the fixed six-forward re-noise sampler, T2VA/I2VA/L2VA/FL2VA/Ref2VA pipeline support, an example, docs, and CPU unit tests. GPU end-to-end validation is still pending. This change does not add PDD, VSA, or new kernel paths.
This commit is contained in:
aryan5v
@@ -0,0 +1,180 @@
|
||||
# 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. Kernel/VSA work is separate.
|
||||
|
||||
Examples:
|
||||
python examples/inference/basic/basic_pika_h3_flex.py --task t2va \
|
||||
--prompt "a paper boat drifting down a rainy street"
|
||||
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
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
ComponentConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference
|
||||
|
||||
# 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")
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--model-path", default="Pika-Labs/Pika-H3-Flex")
|
||||
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("--height", type=int, default=768)
|
||||
parser.add_argument("--width", type=int, default=1344)
|
||||
parser.add_argument("--num-frames", type=int, default=124)
|
||||
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("--low-vram", action="store_true", help="stream the DiT layerwise")
|
||||
parser.add_argument("--output", default="outputs/pika_h3_flex")
|
||||
args = parser.parse_args()
|
||||
|
||||
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")
|
||||
return args
|
||||
|
||||
|
||||
def _inputs(args: argparse.Namespace) -> InputConfig:
|
||||
if args.task == "ref2va":
|
||||
return InputConfig(references=args.reference)
|
||||
return InputConfig(image_path=args.image, last_image=args.last_image)
|
||||
|
||||
|
||||
def _pipeline(args: argparse.Namespace) -> PipelineSelection:
|
||||
components = ComponentConfig()
|
||||
if args.task == "ref2va":
|
||||
components.override_pipeline_cls_name = "PikaH3FlexRef2VAModularPipeline"
|
||||
return PipelineSelection(
|
||||
workload_type="t2v" if args.task == "t2va" else "i2v",
|
||||
preset=f"pika_h3_flex_{args.task}",
|
||||
components=components,
|
||||
vae_tiling=True,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.lossless_accel:
|
||||
for key, value in LOSSLESS_ACCEL_ENV.items():
|
||||
os.environ.setdefault(key, value)
|
||||
|
||||
output_path = Path(args.output) / f"pika_h3_flex_{args.task}.mp4"
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
offload = OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=args.low_vram,
|
||||
text_encoder=True,
|
||||
vae=True,
|
||||
pin_cpu_memory=not args.low_vram,
|
||||
lazy_module_load=True,
|
||||
)
|
||||
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),
|
||||
),
|
||||
pipeline=_pipeline(args),
|
||||
))
|
||||
try:
|
||||
result = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt="",
|
||||
inputs=_inputs(args),
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.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()
|
||||
Reference in New Issue
Block a user