diff --git a/docs/inference/pika_h3_flex.md b/docs/inference/pika_h3_flex.md new file mode 100644 index 000000000..da4dc55fa --- /dev/null +++ b/docs/inference/pika_h3_flex.md @@ -0,0 +1,100 @@ +# Pika H3 Flex + +Pika H3 Flex is a post-trained MiniMax-H3 model that keeps H3's packed +video/audio transformer, Qwen3-VL conditioner, and video/audio VAEs. Pika ships +two rank-128 PEFT adapters instead of a second base checkpoint: + +| Adapter | Trunk | Tasks | +|---|---|---| +| `fl2va/adapter.safetensors` | `transformer` | T2VA, I2VA, L2VA, FL2VA | +| `ref2va/adapter.safetensors` | `transformer_ref` | Ref2VA | + +FastVideo detects `Pika-Labs/Pika-H3-Flex`, downloads only the selected task's +adapter, and applies it as an ordinary LoRA. The released adapter is selected +automatically; pass `lora_path` only when overriding it with a compatible +Pika adapter. + +## Sampling contract + +Pika's sampler is fixed, not a tunable step count: + +* 6 transformer forwards over clean-time rungs + `1000, 750, 500, 250, 100, 40, 0`; +* video shift `12.0`, audio shift `3.0`; +* no CFG, `guidance_scale=1.0`, empty negative prompt; +* each step re-noises the predicted clean sample with fresh Gaussian noise from + a generator seeded `seed + 1`. + +`num_inference_steps` must therefore be `6`; other values are rejected. This is +not a PDD checkpoint and it has no trained VSA gates, so PDD fused blocks and +VSA sparsity are not used. + +## Tasks + +```bash +# text to video + audio +python examples/inference/basic/basic_pika_h3_flex.py --task t2va \ + --prompt "a paper boat drifting down a rainy street" + +# first frame, last frame, or both +python examples/inference/basic/basic_pika_h3_flex.py --task i2va \ + --image first.png --prompt "the camera slowly pushes in" +python examples/inference/basic/basic_pika_h3_flex.py --task fl2va \ + --image first.png --last-image last.png --prompt "the flower opens" + +# ordered references, including reference video and audio +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" +``` + +Ref2VA accepts up to 12 references: 9 images, 3 videos, and 3 audio. Each video +and audio reference is 2–15 s; video/audio totals are capped at 15 s per +modality. A video reference needs at least 22 decoded frames at 24 fps. At +least one image or video is required, and the Ref2VA pipeline must be selected +with `PikaH3FlexRef2VAModularPipeline` (the example does this). + +## Canvas and frames + +Pika uses an exact 32-pixel bucket lattice. A request that matches a bucket is +kept exactly; otherwise FastVideo picks the tier from the requested short edge +(`<= 480` → p480, else p768) and snaps to the nearest aspect bucket in that tier. + +| Tier | Short edge | Max pixels | Reference-image short edge | +|---|---|---|---| +| p480 | 480 | 832×480 | 1024 | +| p768 | 768 | 1344×768 | 2048 | + +Frame counts follow Pika's `17n + 5` grid at 24 fps (5–15 s, 120–362 aligned +frames). The default 16:9 p768 canvas is 1344×768 and 124 frames. + +Pika center-crops both keyframes to the bucket. MiniMax-H3 stretches its first +frame instead; do not reuse that fit for Pika. + +## Training-free acceleration + +The 6 forwards already use FastVideo's normal inference path. Available +lossless switches: + +* `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS=all` — exact fused H3 kernels; +* `FASTVIDEO_H3_REF2VA_MEMO_ENTRIES=16` — cache reference conditioning + (Ref2VA); +* `FASTVIDEO_H3_VAE_TILE_PARALLEL=1` — parallel VAE tiles; +* `FASTVIDEO_ULYSSES_A2A=auto` — Ulysses all-to-all for multi-GPU sequence + parallelism; +* `OffloadConfig(lazy_module_load=True)` and `dit_layerwise=True` for smaller + GPUs; +* `CompileConfig(enabled=True)` for `torch.compile`; +* the standard FlashAttention/SDPA backend selection. + +`--lossless-accel` in the example enables the first four before the generator +starts. Kernel work beyond the existing exact kernels is intentionally not part +of this integration. + +## Hardware + +The selected H3 trunk plus Qwen3-VL does not fit comfortably on a 24 GB card +without offload. A single 80 GB GPU can run with text-encoder offload; use +layerwise DiT offload or multiple GPUs for smaller devices. Ref2VA encodes +reference images at up to a 2048 short edge, so its conditioning memory is +higher than T2VA's. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 722e5cf28..3aa8e43d2 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -68,6 +68,7 @@ column links a runnable script in `examples/inference/basic/` where one exists. | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`
`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`
`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`
`FastVideo/Matrix-Game-2.0-Base-Diffusers`
`FastVideo/Matrix-Game-2.0-GTA-Diffusers`
`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`
`mignonjia/mg_longtuning_distilled_zelda`
`mignonjia/mg_sf_distilled_zelda_1k_steps`
`mignonjia/mg_sf_distilled_zelda`
`mignonjia/mg_causal_zelda`
`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) | | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) | | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)
[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)
[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) | +| pika_h3_flex | `Pika-Labs/Pika-H3-Flex` | T2V, I2V | [T2VA/I2VA/L2VA/FL2VA/Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pika_h3_flex.py) | | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) | | qwen_image21 | `Qwen/Qwen-Image-2.1` | T2I, I2I | [basic_qwen_image21.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_qwen_image21.py) | | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) | diff --git a/examples/inference/basic/basic_pika_h3_flex.py b/examples/inference/basic/basic_pika_h3_flex.py new file mode 100644 index 000000000..0894bdd51 --- /dev/null +++ b/examples/inference/basic/basic_pika_h3_flex.py @@ -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() diff --git a/fastvideo/api/pika_h3_flex.py b/fastvideo/api/pika_h3_flex.py new file mode 100644 index 000000000..2e77eba5b --- /dev/null +++ b/fastvideo/api/pika_h3_flex.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling defaults for Pika H3 Flex.""" + +from dataclasses import dataclass + +from fastvideo.api.sampling_param import SamplingParam + + +@dataclass +class PikaH3FlexSamplingParam(SamplingParam): + """Six re-noise forwards, no guidance, and the 768p bucket by default.""" + + num_inference_steps: int = 6 + guidance_scale: float = 1.0 + batch_cfg: bool = False + negative_prompt: str = "" + height: int = 768 + width: int = 1344 + num_frames: int = 124 + fps: int = 24 + seed: int = 0 + + +__all__ = ["PikaH3FlexSamplingParam"] diff --git a/fastvideo/configs/models/dits/pika_h3_flex.py b/fastvideo/configs/models/dits/pika_h3_flex.py new file mode 100644 index 000000000..ebb9db1c1 --- /dev/null +++ b/fastvideo/configs/models/dits/pika_h3_flex.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex transformer configuration. + +The released Pika adapters do not change MiniMax-H3's transformer: the +checkpoint keeps the same packed video/audio trunk, Qwen3-VL conditioner, and +VAEs, and applies a rank-128 PEFT adapter at inference. This subclass exists so +the Pika pipeline names its own component config without duplicating the H3 +architecture contract. +""" + +from dataclasses import dataclass + +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config + + +@dataclass +class PikaH3FlexDiTConfig(MiniMaxH3Config): + """The released MiniMax-H3 transformer, sampled through Pika's adapters.""" + + +__all__ = ["PikaH3FlexDiTConfig"] diff --git a/fastvideo/configs/pipelines/pika_h3_flex.py b/fastvideo/configs/pipelines/pika_h3_flex.py new file mode 100644 index 000000000..d0d674e19 --- /dev/null +++ b/fastvideo/configs/pipelines/pika_h3_flex.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline configuration for Pika H3 Flex. + +Pika H3 Flex is a post-trained MiniMax-H3: the transformer, Qwen3-VL +conditioner, and video/audio VAEs are the released H3 components, and a rank-128 +adapter is merged at load time. What changes is the sampler and the geometry +contract: + +* six transformer forwards over the fixed clean-time knots + ``1000, 750, 500, 250, 100, 40``, with a fresh-noise re-noise step between + forwards (the terminal ``0`` knot is the output, never a forward); +* video and audio clock shifts of 12 and 3; +* 480p and 768p canvas buckets instead of a free-form canvas; +* center-cropped keyframes and tier-dependent reference-image sizing. + +Everything else, including CFG being off and the 24 fps frame grid, is inherited +from MiniMax-H3. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from fastvideo.configs.models.dits.pika_h3_flex import PikaH3FlexDiTConfig +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + +if TYPE_CHECKING: + from fastvideo.api.sampling_param import SamplingParam + from fastvideo.api.schema import GenerationRequest + +# The released re-noise ladder. The terminal 0 is a sigma-grid point, not a +# transformer forward: six forwards, seven sigmas. +PIKA_DENOISING_STEPS: tuple[int, ...] = (1000, 750, 500, 250, 100, 40) +PIKA_VIDEO_SCHEDULER_SHIFT = 12.0 +PIKA_AUDIO_SCHEDULER_SHIFT = 3.0 +PIKA_RESOLUTION_TIERS: tuple[str, ...] = ("p480", "p768") +PIKA_DEFAULT_RESOLUTION_TIER = "p768" + + +@dataclass +class PikaH3FlexPipelineConfig(MiniMaxH3PipelineConfig): + """MiniMax-H3 components under Pika's adapter, sampler, and bucket contract.""" + + dit_config: PikaH3FlexDiTConfig = field(default_factory=PikaH3FlexDiTConfig) + # Pika's CLI calls this ``--resolution``; the API has no dedicated tier + # field, so it is a pipeline config knob rather than a per-request one. + pika_resolution_tier: str = PIKA_DEFAULT_RESOLUTION_TIER + pika_denoising_steps: tuple[int, ...] = PIKA_DENOISING_STEPS + + def check_pipeline_config(self) -> None: + super().check_pipeline_config() + if self.pika_resolution_tier not in PIKA_RESOLUTION_TIERS: + raise ValueError(f"Pika H3 Flex supports resolution tiers {list(PIKA_RESOLUTION_TIERS)}; got " + f"{self.pika_resolution_tier!r}.") + if tuple(self.pika_denoising_steps) != PIKA_DENOISING_STEPS: + raise ValueError(f"Pika H3 Flex's adapter contract fixes the denoising ladder at " + f"{list(PIKA_DENOISING_STEPS)}; got {list(self.pika_denoising_steps)}.") + if self.dmd_denoising_steps is not None: + raise ValueError("Pika H3 Flex uses its own six-forward re-noise schedule; dmd_denoising_steps must " + "remain unset.") + if self.pdd_step_indices is not None: + raise ValueError("Pika H3 Flex is not a FastH3 PDD student; pdd_step_indices must remain unset.") + + def apply_request_constraints(self, request: GenerationRequest, sampling_param: SamplingParam) -> SamplingParam: + """Fix ``num_inference_steps`` to the adapter's six-forward ladder.""" + from fastvideo.api.compat import explicit_request_updates + + steps = len(PIKA_DENOISING_STEPS) + requested = explicit_request_updates(request).get("num_inference_steps", steps) + if requested != steps: + raise ValueError(f"Pika H3 Flex runs exactly {steps} transformer forwards; the request sets " + f"num_inference_steps={requested}. The six-step ladder is part of the adapter " + f"contract, not a quality/speed knob.") + sampling_param.num_inference_steps = steps + return sampling_param + + +__all__ = [ + "PIKA_AUDIO_SCHEDULER_SHIFT", + "PIKA_DEFAULT_RESOLUTION_TIER", + "PIKA_DENOISING_STEPS", + "PIKA_RESOLUTION_TIERS", + "PIKA_VIDEO_SCHEDULER_SHIFT", + "PikaH3FlexPipelineConfig", +] diff --git a/fastvideo/models/schedulers/scheduling_minimax_h3.py b/fastvideo/models/schedulers/scheduling_minimax_h3.py index 647075caa..ef901f32e 100644 --- a/fastvideo/models/schedulers/scheduling_minimax_h3.py +++ b/fastvideo/models/schedulers/scheduling_minimax_h3.py @@ -21,13 +21,31 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin): order = 1 @register_to_config - def __init__(self, shift: float = 12.0) -> None: + def __init__( + self, + shift: float = 12.0, + stochastic_sampling: bool = False, + shift_provided_sigmas: bool = False, + ) -> None: + """H3's rectified-flow scheduler, with Pika's re-noise step as an option. + + ``stochastic_sampling`` switches the Euler update to the re-noise + contract: draw fresh Gaussian noise at the next sigma and form + ``(1 - sigma_next) * x0 + sigma_next * noise``. The terminal sigma + always returns ``x0`` without drawing noise. ``shift_provided_sigmas`` + records that explicit sigmas are already shifted (Pika's release uses + it); FastVideo always treats explicit ``set_timesteps(sigmas=...)`` + values that way, so it is informational. + """ if shift <= 0: raise ValueError(f"`shift` must be positive, got {shift}.") self.num_inference_steps: int | None = None self.sigmas: torch.Tensor | None = None self.timesteps: torch.Tensor | None = None self._shift = float(shift) + self._stochastic_sampling = bool(stochastic_sampling) + self._shift_provided_sigmas = bool(shift_provided_sigmas) + self._noise_generator: torch.Generator | None = None self._step_index: int | None = None self._begin_index: int | None = None @@ -35,6 +53,18 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin): def shift(self) -> float: return self._shift + @property + def stochastic_sampling(self) -> bool: + return self._stochastic_sampling + + @property + def shift_provided_sigmas(self) -> bool: + return self._shift_provided_sigmas + + def set_noise_generator(self, generator: torch.Generator | None) -> None: + """Set the generator used by stochastic (re-noise) steps.""" + self._noise_generator = generator + @property def step_index(self) -> int | None: return self._step_index @@ -127,11 +157,29 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin): denoised = sample + sigma_from_timestep * model_output compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype - sigma = self.sigmas[self._step_index].to(device=sample.device, dtype=compute_dtype) sigma_next = self.sigmas[self._step_index + 1].to(device=sample.device, dtype=compute_dtype) - ratio = sigma_next / sigma - prev_sample = ratio * sample.to(compute_dtype) + (1.0 - ratio) * denoised.to(compute_dtype) - prev_sample = prev_sample.to(sample.dtype) + if self._stochastic_sampling: + if self._noise_generator is None: + raise ValueError("Stochastic MiniMax-H3 sampling needs a noise generator; call " + "`set_noise_generator` before `step`.") + if float(sigma_next) == 0.0: + # The terminal ladder point is the generated sample. Pika's + # sampler returns x0 here instead of drawing a final noise. + prev_sample = denoised.to(sample.dtype) + else: + noise = torch.randn( + sample.shape, + generator=self._noise_generator, + device=sample.device, + dtype=torch.float32, + ) + prev_sample = ((1.0 - sigma_next) * denoised.to(compute_dtype) + + sigma_next * noise).to(sample.dtype) + else: + sigma = self.sigmas[self._step_index].to(device=sample.device, dtype=compute_dtype) + ratio = sigma_next / sigma + prev_sample = ratio * sample.to(compute_dtype) + (1.0 - ratio) * denoised.to(compute_dtype) + prev_sample = prev_sample.to(sample.dtype) self._step_index += 1 if not return_dict: diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index 9cee168e8..797a89509 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -218,6 +218,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): pipeline_config_cls: type[MiniMaxH3PipelineConfig] = MiniMaxH3PipelineConfig _ref2va_default = False + # Family subclasses override these to swap geometry and sampler stages + # without copying the whole stage assembly. + input_preparation_stage_cls: type[MiniMaxH3InputPreparationStage] = MiniMaxH3InputPreparationStage + denoising_stage_cls: type[MiniMaxH3DenoisingStage] = MiniMaxH3DenoisingStage _required_config_modules = [ "text_encoder", "tokenizer", @@ -686,7 +690,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): def _add_input_stage(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None: self.add_stage( "input_preparation_stage", - MiniMaxH3InputPreparationStage( + self.input_preparation_stage_cls( vae=self._input_video_geometry(fastvideo_args), audio_vae=self._input_audio_vae(fastvideo_args, ref2va=ref2va), ref2va=ref2va, @@ -734,7 +738,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): ) self.add_stage( "denoising_stage", - MiniMaxH3DenoisingStage( + self.denoising_stage_cls( transformer=transformer, scheduler=scheduler, audio_scheduler=audio_scheduler, diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index 1b947dfed..5c91a7207 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -236,20 +236,33 @@ def build_packed_batch_sequence(documents: tuple[MiniMaxH3PackedLayout, ...], ) ) -def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: +def resolve_canvas_size( + aspect_width: float, + aspect_height: float, + *, + short_edge: int = MINIMAX_H3_SHORT_EDGE, + max_pixels: int = MINIMAX_H3_MAX_PIXELS, +) -> tuple[int, int]: + """Resolve a canvas from an aspect ratio and a tier's short-edge/area rules. + + The defaults are the released MiniMax-H3 768p rule. Pika H3 Flex passes its + own 480p/768p tier values through the keyword arguments. + """ if aspect_width <= 0 or aspect_height <= 0: raise ValueError(f"The aspect ratio must be positive, got {aspect_width}:{aspect_height}.") + if short_edge <= 0 or max_pixels <= 0: + raise ValueError(f"Canvas short_edge and max_pixels must be positive, got {short_edge}, {max_pixels}.") ratio = aspect_width / aspect_height if not MINIMAX_H3_MIN_ASPECT_RATIO <= ratio <= MINIMAX_H3_MAX_ASPECT_RATIO: raise ValueError(f"MiniMax-H3 supports aspect ratios from 1:4 to 4:1, got {aspect_width}:{aspect_height}.") if ratio >= 1: - width, height = MINIMAX_H3_SHORT_EDGE * ratio, float(MINIMAX_H3_SHORT_EDGE) + width, height = short_edge * ratio, float(short_edge) else: - width, height = float(MINIMAX_H3_SHORT_EDGE), MINIMAX_H3_SHORT_EDGE / ratio + width, height = float(short_edge), short_edge / ratio area = width * height - if area > MINIMAX_H3_MAX_PIXELS: - scale = (MINIMAX_H3_MAX_PIXELS / area)**0.5 + if area > max_pixels: + scale = (max_pixels / area)**0.5 width, height = width * scale, height * scale multiple = MINIMAX_H3_CANVAS_MULTIPLE return max(multiple, round(height / multiple) * multiple), max(multiple, round(width / multiple) * multiple) diff --git a/fastvideo/pipelines/basic/minimax_h3/reference.py b/fastvideo/pipelines/basic/minimax_h3/reference.py index 712cdd6bf..965d3327d 100644 --- a/fastvideo/pipelines/basic/minimax_h3/reference.py +++ b/fastvideo/pipelines/basic/minimax_h3/reference.py @@ -22,6 +22,8 @@ from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_FRAMES_PER_CHUNK, MINIMAX_H3_LATENTS_PER_CHUNK, MINIMAX_H3_MAX_DURATION, + MINIMAX_H3_MAX_PIXELS, + MINIMAX_H3_SHORT_EDGE, resolve_canvas_size, ) @@ -180,13 +182,20 @@ def decode_reference_video(source: str | os.PathLike[str]) -> tuple[np.ndarray, return result, float(rate), soundtrack -def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: - """Resolve the released 2048-short-edge reference-image canvas.""" +def resolve_reference_image_size( + width: int, + height: int, + *, + short_edge: int = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE, +) -> tuple[int, int]: + """Resolve the released reference-image canvas at ``short_edge``.""" if width <= 0 or height <= 0: raise ValueError(f"A reference image must have a positive size, got {width}x{height}.") if width > 4 * height or height > 4 * width: raise ValueError(f"A reference image must be within 1:4 and 4:1, got {width}x{height}.") - scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + if short_edge <= 0: + raise ValueError(f"A reference-image short edge must be positive, got {short_edge}.") + scale = short_edge / min(width, height) multiple = MINIMAX_H3_CANVAS_MULTIPLE return ( max(multiple, @@ -231,12 +240,23 @@ def resample_reference_frames(frames: np.ndarray, fps: float) -> np.ndarray: return np.repeat(frames, repeats, axis=0) -def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray: +def prepare_reference_frames( + frames: np.ndarray, + num_frames: int, + *, + short_edge: int = MINIMAX_H3_SHORT_EDGE, + max_pixels: int = MINIMAX_H3_MAX_PIXELS, +) -> np.ndarray: """Trim a reference video to the request and resize every RGB frame.""" if frames.ndim != 4 or frames.shape[-1] != 3 or frames.shape[0] == 0: raise ValueError(f"A reference video must be non-empty RGB frames, got {tuple(frames.shape)}.") frames = frames[:num_frames] - height, width = resolve_canvas_size(frames.shape[2], frames.shape[1]) + height, width = resolve_canvas_size( + frames.shape[2], + frames.shape[1], + short_edge=short_edge, + max_pixels=max_pixels, + ) if frames.shape[1:3] == (height, width): return frames return np.stack( @@ -319,8 +339,17 @@ def prepare_reference( reference: MiniMaxH3Reference, num_frames: int, target_sample_rate: int, + *, + reference_short_edge: int = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE, + reference_canvas_short_edge: int = MINIMAX_H3_SHORT_EDGE, + reference_canvas_max_pixels: int = MINIMAX_H3_MAX_PIXELS, ) -> MiniMaxH3PreparedReference: - """Decode and normalize one deferred Ref2VA medium.""" + """Decode and normalize one deferred Ref2VA medium. + + The keyword arguments let a family use its own tier's reference-image short + edge and reference-video canvas rule. MiniMax-H3's released values are the + defaults; Pika H3 Flex passes its 480p/768p tier values. + """ prepared = MiniMaxH3PreparedReference(media_type=reference.media_type) if reference.media_type == "image": source = reference.source @@ -331,7 +360,7 @@ def prepare_reference( raise ValueError(f"An image reference must be RGB, got {tuple(pixels.shape)}.") image = Image.fromarray(pixels) image = ImageOps.exif_transpose(image).convert("RGB") - height, width = resolve_reference_image_size(*image.size) + height, width = resolve_reference_image_size(*image.size, short_edge=reference_short_edge) prepared.image = prepare_reference_image(image, height, width) return prepared @@ -346,7 +375,12 @@ def prepare_reference( frames = reference_media_to_uint8(reference.source) decoded_fps = float(MINIMAX_H3_FPS) fps = float(reference.fps if reference.fps is not None else decoded_fps) - prepared.frames = prepare_reference_frames(resample_reference_frames(frames, fps), num_frames) + prepared.frames = prepare_reference_frames( + resample_reference_frames(frames, fps), + num_frames, + short_edge=reference_canvas_short_edge, + max_pixels=reference_canvas_max_pixels, + ) if reference.soundtrack is not None: waveform, source_rate = _resolve_audio_source(reference.soundtrack) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 11e286823..d8333997f 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -133,14 +133,21 @@ class MiniMaxH3DenoisingStage(PipelineStage): self.scheduler = scheduler self.audio_scheduler = audio_scheduler - def _set_dmd_schedule(self, steps: list[int], grid_points: int, device: torch.device) -> None: - """Run the trained rungs, shifting the shared noise clock once per modality.""" + def _set_shifted_schedule( + self, + steps: list[int], + grid_points: int, + device: torch.device, + *, + label: str = "MiniMax-H3 DMD", + ) -> None: + """Build one modality-shifted sigma grid from explicit clean-time rungs.""" if (not steps or any(type(step) is not int or not 0 < step <= 1000 for step in steps) or any(left <= right for left, right in zip(steps, steps[1:], strict=False))): - raise ValueError("MiniMax-H3 DMD rungs must be strictly decreasing integers in (0, 1000].") + raise ValueError(f"{label} rungs must be strictly decreasing integers in (0, 1000].") if grid_points != len(steps) + 1: - raise ValueError("MiniMax-H3 num_inference_steps counts sigma-grid points: " - f"{len(steps)} DMD forwards require {len(steps) + 1} grid points, got {grid_points}.") + raise ValueError(f"{label} num_inference_steps counts sigma-grid points: " + f"{len(steps)} forwards require {len(steps) + 1} grid points, got {grid_points}.") base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32) for scheduler in (self.scheduler, self.audio_scheduler): shift = float(scheduler.shift) @@ -149,6 +156,38 @@ class MiniMaxH3DenoisingStage(PipelineStage): # clean time (1 - sigma); passing integer rungs to step() is wrong. scheduler.set_timesteps(sigmas=sigmas, device=device) + def _set_dmd_schedule(self, steps: list[int], grid_points: int, device: torch.device) -> None: + """Run the trained rungs, shifting the shared noise clock once per modality.""" + self._set_shifted_schedule(steps, grid_points, device, label="MiniMax-H3 DMD") + + def _configure_schedules( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + device: torch.device, + ) -> PDDSamplingPlan | None: + """Install this checkpoint's sigma schedules and return a PDD plan, if any. + + Families with a fixed sampler (Pika H3 Flex) override this hook; the + default keeps MiniMax-H3's PDD and DMD checkpoint contracts. + """ + dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps + pdd_plan = self._pdd_sampling_plan(fastvideo_args, device) + if pdd_plan is not None: + # One runtime step fuses one block of heads into their weighted + # mean velocity; the ordinary Euler step over the block's two node + # sigmas then applies the block's total increment per modality. + self.scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["video"].to(torch.float32), device=device) + self.audio_scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["audio"].to(torch.float32), device=device) + logger.info("MiniMax-H3 PDD denoising: %d fused block(s) %s over a %d-interval fine grid.", + pdd_plan.num_steps, pdd_plan.block_sizes, pdd_plan.pdd_steps) + elif dmd_steps is None: + self.scheduler.set_timesteps(batch.num_inference_steps, device=device) + self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) + else: + self._set_dmd_schedule(dmd_steps, batch.num_inference_steps, device) + return pdd_plan + def _pdd_sampling_plan(self, fastvideo_args: FastVideoArgs, device: torch.device) -> PDDSamplingPlan | None: """The fused-block plan of a PDD checkpoint, or None for any other checkpoint. @@ -207,21 +246,7 @@ class MiniMaxH3DenoisingStage(PipelineStage): and not fastvideo_args.use_fsdp_inference) device = get_local_torch_device() - dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps - pdd_plan = self._pdd_sampling_plan(fastvideo_args, device) - if pdd_plan is not None: - # One runtime step fuses one block of heads into their weighted - # mean velocity; the ordinary Euler step over the block's two node - # sigmas then applies the block's total increment per modality. - self.scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["video"].to(torch.float32), device=device) - self.audio_scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["audio"].to(torch.float32), device=device) - logger.info("MiniMax-H3 PDD denoising: %d fused block(s) %s over a %d-interval fine grid.", - pdd_plan.num_steps, pdd_plan.block_sizes, pdd_plan.pdd_steps) - elif dmd_steps is None: - self.scheduler.set_timesteps(batch.num_inference_steps, device=device) - self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) - else: - self._set_dmd_schedule(dmd_steps, batch.num_inference_steps, device) + pdd_plan = self._configure_schedules(batch, fastvideo_args, device) video_timesteps = self.scheduler.timesteps audio_timesteps = self.audio_scheduler.timesteps if video_timesteps is None or audio_timesteps is None: diff --git a/fastvideo/pipelines/basic/pika_h3_flex/__init__.py b/fastvideo/pipelines/basic/pika_h3_flex/__init__.py new file mode 100644 index 000000000..074ce47e7 --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.pika_h3_flex.pipeline import ( + EntryClass, + PikaH3FlexModularPipeline, + PikaH3FlexRef2VAModularPipeline, +) +from fastvideo.pipelines.basic.pika_h3_flex.adapter import PikaAdapter + +__all__ = [ + "EntryClass", + "PikaAdapter", + "PikaH3FlexModularPipeline", + "PikaH3FlexRef2VAModularPipeline", +] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/adapter.py b/fastvideo/pipelines/basic/pika_h3_flex/adapter.py new file mode 100644 index 000000000..7177c0d23 --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/adapter.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Locate and validate Pika H3 Flex's released LoRA adapters. + +Pika H3 Flex is MiniMax-H3 plus one of two rank-128 PEFT adapters: ``fl2va`` +for text/first-frame/last-frame generation and ``ref2va`` for ordered +image/video/audio references. The adapter is part of the model contract, not a +user-tunable accelerator, so the pipeline loads it automatically and refuses a +manifest whose rank, trunk, or sampler does not match the release. + +This module stays torch-free; the manifest is a plain JSON document and the +only optional dependency is ``huggingface_hub`` for a repo-id source. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from collections.abc import Mapping + +from fastvideo.configs.pipelines.pika_h3_flex import ( + PIKA_AUDIO_SCHEDULER_SHIFT, + PIKA_DENOISING_STEPS, + PIKA_VIDEO_SCHEDULER_SHIFT, +) + +PIKA_ADAPTER_FORMAT_VERSION = 1 +PIKA_ADAPTER_RANK = 128 +PIKA_ADAPTER_ALPHA = 128 +PIKA_BASE_REPO = "MiniMaxAI/MiniMax-H3" +PIKA_ADAPTER_FILENAME = "adapter.safetensors" +PIKA_MANIFEST_FILENAME = "manifest.json" + +PIKA_SAMPLING_CONTRACT = { + "knots": (*PIKA_DENOISING_STEPS, 0), + "video_shift": PIKA_VIDEO_SCHEDULER_SHIFT, + "audio_shift": PIKA_AUDIO_SCHEDULER_SHIFT, + "method": "renoise", +} +PIKA_TARGET_MODULES = ( + "to_q", + "to_k", + "to_v", + "to_out.0", + "ff.net.0.proj", + "ff.net.2", +) + +TASK_TO_TRUNK = {"fl2va": "transformer", "ref2va": "transformer_ref"} + + +class PikaAdapterError(ValueError): + """The Pika adapter is missing or does not describe the released model.""" + + +@dataclass(frozen=True) +class PikaAdapter: + """One validated released adapter.""" + + task: str + trunk: str + path: str + manifest_path: str + manifest: dict[str, Any] + + @property + def directory(self) -> str: + return str(Path(self.path).parent) + + +def task_for_pipeline(ref2va: bool) -> str: + return "ref2va" if ref2va else "fl2va" + + +def expected_trunk(task: str) -> str: + try: + return TASK_TO_TRUNK[task] + except KeyError as error: + raise PikaAdapterError( + f"Unknown Pika H3 Flex task {task!r}; expected one of {sorted(TASK_TO_TRUNK)}.") from error + + +def _as_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise PikaAdapterError(f"Pika adapter manifest field {field!r} must be an integer, got {value!r}.") + return value + + +def _as_float(value: Any, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise PikaAdapterError(f"Pika adapter manifest field {field!r} must be a number, got {value!r}.") + return float(value) + + +def validate_adapter_manifest( + manifest: Mapping[str, Any], + *, + task: str, + pipeline_trunk: str | None = None, +) -> None: + """Reject a manifest that is not the released Pika adapter for ``task``.""" + trunk = expected_trunk(task) + if pipeline_trunk is not None and pipeline_trunk != trunk: + raise PikaAdapterError(f"Pika H3 Flex task {task!r} loads trunk {trunk!r}, but the pipeline requested " + f"{pipeline_trunk!r}.") + + format_version = _as_int(manifest.get("format_version"), "format_version") + if format_version != PIKA_ADAPTER_FORMAT_VERSION: + raise PikaAdapterError(f"Unsupported Pika adapter manifest format {format_version}; expected " + f"{PIKA_ADAPTER_FORMAT_VERSION}.") + manifest_task = manifest.get("task") + if manifest_task is not None and manifest_task != task: + raise PikaAdapterError(f"The adapter at this path is for task {manifest_task!r}, not {task!r}.") + rank = _as_int(manifest.get("rank"), "rank") + alpha = _as_int(manifest.get("alpha"), "alpha") + if rank != PIKA_ADAPTER_RANK or alpha != PIKA_ADAPTER_ALPHA: + raise PikaAdapterError(f"Pika H3 Flex adapters are rank {PIKA_ADAPTER_RANK} with alpha " + f"{PIKA_ADAPTER_ALPHA}; this manifest says rank {rank}, alpha {alpha}.") + manifest_trunk = manifest.get("trunk") + if manifest_trunk != trunk: + raise PikaAdapterError(f"Pika H3 Flex task {task!r} expects adapter trunk {trunk!r}, got " + f"{manifest_trunk!r}.") + + sampling = manifest.get("sampling") + if not isinstance(sampling, Mapping): + raise PikaAdapterError("Pika adapter manifest is missing its `sampling` contract.") + knots = sampling.get("knots") + if not isinstance(knots, list | tuple) or tuple(knots) != PIKA_SAMPLING_CONTRACT["knots"]: + raise PikaAdapterError(f"Pika adapter sampling knots must be " + f"{list(PIKA_SAMPLING_CONTRACT['knots'])}, got {knots!r}.") + for field in ("video_shift", "audio_shift"): + value = _as_float(sampling.get(field), f"sampling.{field}") + expected = float(PIKA_SAMPLING_CONTRACT[field]) + if abs(value - expected) > 1e-9: + raise PikaAdapterError(f"Pika adapter sampling.{field} must be {expected:g}, got {value:g}.") + if sampling.get("method") != PIKA_SAMPLING_CONTRACT["method"]: + raise PikaAdapterError(f"Pika adapter sampling.method must be " + f"{PIKA_SAMPLING_CONTRACT['method']!r}, got {sampling.get('method')!r}.") + + target_modules = manifest.get("target_modules") + if not isinstance(target_modules, list | tuple) or tuple(target_modules) != PIKA_TARGET_MODULES: + raise PikaAdapterError(f"Pika adapter target_modules must be {list(PIKA_TARGET_MODULES)}, got " + f"{target_modules!r}.") + + base = manifest.get("base") + if not isinstance(base, Mapping) or base.get("repo_id") != PIKA_BASE_REPO: + raise PikaAdapterError(f"Pika adapter base must name {PIKA_BASE_REPO!r}, got {base!r}.") + + +def _load_manifest(path: Path) -> dict[str, Any]: + try: + manifest = json.loads(path.read_text(encoding="utf-8")) + except OSError as error: + raise PikaAdapterError(f"Cannot read Pika adapter manifest {str(path)!r}: {error}") from error + except json.JSONDecodeError as error: + raise PikaAdapterError(f"Pika adapter manifest {str(path)!r} is not valid JSON: {error}") from error + if not isinstance(manifest, dict): + raise PikaAdapterError(f"Pika adapter manifest {str(path)!r} must be a JSON object.") + return manifest + + +def _adapter_in_directory(directory: Path, task: str) -> tuple[Path, Path] | None: + nested = directory / task / PIKA_ADAPTER_FILENAME + if nested.is_file(): + manifest = directory / task / PIKA_MANIFEST_FILENAME + return nested, manifest + flat = directory / PIKA_ADAPTER_FILENAME + if flat.is_file(): + manifest = directory / PIKA_MANIFEST_FILENAME + return flat, manifest + return None + + +def _looks_like_repo_id(source: str) -> bool: + parts = source.split("/") + return len(parts) == 2 and all(parts) and "\\" not in source + + +def _download_adapter(source: str, task: str, revision: str | None) -> tuple[Path, Path]: + try: + from huggingface_hub import hf_hub_download + except ImportError as error: # pragma: no cover - huggingface_hub is a runtime dependency + raise PikaAdapterError("Downloading a Pika adapter by repo id requires huggingface_hub.") from error + try: + adapter = Path(hf_hub_download( + repo_id=source, + filename=f"{task}/{PIKA_ADAPTER_FILENAME}", + revision=revision, + )) + manifest = Path( + hf_hub_download( + repo_id=source, + filename=f"{task}/{PIKA_MANIFEST_FILENAME}", + revision=revision, + )) + except Exception as error: + raise PikaAdapterError(f"Could not download the Pika H3 Flex {task} adapter from {source!r}" + f"{f' at revision {revision!r}' if revision else ''}: {error}") from error + return adapter, manifest + + +def load_adapter(source: str | os.PathLike[str], *, task: str, revision: str | None = None) -> PikaAdapter: + """Resolve ``source`` (adapter file, release directory, task directory, or repo id).""" + trunk = expected_trunk(task) + raw = os.fspath(source) + path = Path(raw).expanduser() + if path.is_file(): + adapter_path, manifest_path = path, path.with_name(PIKA_MANIFEST_FILENAME) + elif path.is_dir(): + found = _adapter_in_directory(path, task) + if found is None: + raise PikaAdapterError(f"{raw!r} has no {task}/{PIKA_ADAPTER_FILENAME} or {PIKA_ADAPTER_FILENAME}.") + adapter_path, manifest_path = found + elif _looks_like_repo_id(raw): + adapter_path, manifest_path = _download_adapter(raw, task, revision) + else: + raise PikaAdapterError(f"Pika adapter path {raw!r} does not exist and is not a Hub repo id.") + + if not adapter_path.is_file(): + raise PikaAdapterError(f"Pika adapter weights are missing at {str(adapter_path)!r}.") + if not manifest_path.is_file(): + raise PikaAdapterError(f"Pika adapter manifest is missing at {str(manifest_path)!r}; a bare " + f"{PIKA_ADAPTER_FILENAME} cannot be checked against the release contract.") + manifest = _load_manifest(manifest_path) + validate_adapter_manifest(manifest, task=task, pipeline_trunk=trunk) + return PikaAdapter( + task=task, + trunk=trunk, + path=str(adapter_path), + manifest_path=str(manifest_path), + manifest=manifest, + ) + + +def adapter_hf_patterns(task: str | None = None) -> list[str]: + """The adapter files the pipeline must include in a partial Hub snapshot.""" + tasks = tuple(TASK_TO_TRUNK) if task is None else (task, ) + patterns: list[str] = [] + for item in tasks: + expected_trunk(item) + patterns.extend((f"{item}/{PIKA_ADAPTER_FILENAME}", f"{item}/{PIKA_MANIFEST_FILENAME}")) + return patterns + + +__all__ = [ + "PIKA_ADAPTER_ALPHA", + "PIKA_ADAPTER_FILENAME", + "PIKA_ADAPTER_FORMAT_VERSION", + "PIKA_ADAPTER_RANK", + "PIKA_BASE_REPO", + "PIKA_MANIFEST_FILENAME", + "PIKA_SAMPLING_CONTRACT", + "PIKA_TARGET_MODULES", + "PikaAdapter", + "PikaAdapterError", + "TASK_TO_TRUNK", + "adapter_hf_patterns", + "expected_trunk", + "load_adapter", + "task_for_pipeline", + "validate_adapter_manifest", +] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/buckets.py b/fastvideo/pipelines/basic/pika_h3_flex/buckets.py new file mode 100644 index 000000000..ef9891e6c --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/buckets.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex resolution buckets and frame grid. + +Pika H3 Flex keeps MiniMax-H3's canvas rule: sides sit on a 32-pixel lattice, +a fixed short edge bounds the bucket, and an area cap keeps the token count +finite. The open-weight release samples at two of the released tiers, 480p and +768p; the hosted API's larger tiers are intentionally not exposed here because +the local checkpoints were not trained for them. + +A request selects a bucket, not a free-form canvas. ``height`` and ``width`` +are therefore aspect hints: an exact bucket is honoured, anything else is +snapped to the nearest aspect family of the tier implied by the requested +short edge. Keeping this arithmetic in a torch-free module lets the geometry +contract be tested on a CPU-only host. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from collections.abc import Iterable + +PIKA_FPS = 24 +PIKA_CANVAS_MULTIPLE = 32 +PIKA_MIN_ASPECT_RATIO = 0.25 +PIKA_MAX_ASPECT_RATIO = 4.0 + +# Landscape families first, then square, then the mirrored portraits. +PIKA_ASPECTS: tuple[str, ...] = ( + "21:9", + "16:9", + "3:2", + "4:3", + "5:4", + "1:1", + "4:5", + "3:4", + "2:3", + "9:16", + "9:21", +) + +# 17n + 5 from 2.3 s to 15.1 s at 24 fps. Open-weight requests are limited to +# the documented 5-15 s range; the two smaller grid points remain part of the +# released grid but are rejected by the request validator. +PIKA_FRAME_GRID: tuple[int, ...] = tuple(17 * n + 5 for n in range(3, 22)) +PIKA_MIN_DURATION = 5.0 +PIKA_MAX_DURATION = 15.0 +PIKA_MIN_ALIGNED_FRAMES = 120 +PIKA_MAX_ALIGNED_FRAMES = 362 + + +@dataclass(frozen=True) +class PikaTier: + """One released resolution tier.""" + + name: str + short_edge: int + max_pixels: int + reference_short_edge: int + + +PIKA_TIERS: dict[str, PikaTier] = { + "p480": PikaTier("p480", short_edge=480, max_pixels=832 * 480, reference_short_edge=1024), + "p768": PikaTier("p768", short_edge=768, max_pixels=1344 * 768, reference_short_edge=2048), +} +DEFAULT_TIER = "p768" +DEFAULT_ASPECT: tuple[int, int] = (16, 9) + +_TIER_ALIASES = { + "480p": "p480", + "p480": "p480", + "480": "p480", + "768p": "p768", + "p768": "p768", + "768": "p768", +} + + +def resolve_tier(value: str | None) -> str: + """Normalize ``480p`` / ``p480`` / ``768p`` / ``p768`` to a released tier name.""" + if value is None: + return DEFAULT_TIER + key = str(value).strip().lower() + tier = _TIER_ALIASES.get(key) + if tier is None: + raise ValueError(f"Pika H3 Flex only samples the released 480p and 768p tiers; got {value!r}.") + return tier + + +def tier_rule(tier: str | None) -> PikaTier: + name = resolve_tier(tier) + return PIKA_TIERS[name] + + +def _parse_aspect(aspect: str | tuple[int, int] | tuple[float, float]) -> tuple[float, float]: + if isinstance(aspect, str): + parts = aspect.split(":") + if len(parts) != 2: + raise ValueError(f"An aspect ratio must be W:H, got {aspect!r}.") + width, height = float(parts[0]), float(parts[1]) + elif isinstance(aspect, tuple) and len(aspect) == 2: + width, height = float(aspect[0]), float(aspect[1]) + else: + raise TypeError(f"An aspect ratio must be a 'W:H' string or a (width, height) pair, got {aspect!r}.") + if not (math.isfinite(width) and math.isfinite(height)) or width <= 0 or height <= 0: + raise ValueError(f"An aspect ratio must be finite and positive, got {aspect!r}.") + return width, height + + +def aspect_ratio(width: float, height: float) -> float: + width, height = _parse_aspect((width, height)) + return width / height + + +def canvas_for_aspect(aspect: str | tuple[float, float], tier: str | None = None) -> tuple[int, int]: + """``(height, width)`` of one tier's bucket for an aspect family.""" + rule = tier_rule(tier) + aspect_width, aspect_height = _parse_aspect(aspect) + ratio = aspect_width / aspect_height + if not PIKA_MIN_ASPECT_RATIO <= ratio <= PIKA_MAX_ASPECT_RATIO: + raise ValueError(f"Pika H3 Flex supports aspect ratios from 1:4 to 4:1, got {aspect!r}.") + + if ratio >= 1: + width, height = rule.short_edge * ratio, float(rule.short_edge) + else: + width, height = float(rule.short_edge), rule.short_edge / ratio + area = width * height + if area > rule.max_pixels: + scale = (rule.max_pixels / area)**0.5 + width, height = width * scale, height * scale + multiple = PIKA_CANVAS_MULTIPLE + return ( + max(multiple, + round(height / multiple) * multiple), + max(multiple, + round(width / multiple) * multiple), + ) + + +def aspect_families(tier: str | None = None) -> tuple[str, ...]: + del tier + return PIKA_ASPECTS + + +def bucket_table(tier: str | None = None) -> tuple[tuple[int, int], ...]: + """All ``(width, height)`` buckets of one tier, in :data:`PIKA_ASPECTS` order.""" + return tuple((canvas_for_aspect(aspect, tier)[1], canvas_for_aspect(aspect, tier)[0]) for aspect in PIKA_ASPECTS) + + +def iter_buckets() -> Iterable[tuple[str, str, int, int]]: + """Yield ``(tier, aspect, width, height)`` for every released bucket.""" + for tier in PIKA_TIERS: + for aspect in PIKA_ASPECTS: + height, width = canvas_for_aspect(aspect, tier) + yield tier, aspect, width, height + + +def exact_bucket(height: int, width: int) -> tuple[str, str] | None: + """The ``(tier, aspect)`` whose bucket is exactly ``height x width``, if any.""" + for tier, aspect, bucket_width, bucket_height in iter_buckets(): + if (bucket_height, bucket_width) == (int(height), int(width)): + return tier, aspect + return None + + +def tier_for_canvas(height: int, width: int) -> str: + """The smallest released tier whose short edge covers the canvas.""" + if height <= 0 or width <= 0: + raise ValueError(f"A canvas must be positive, got {height}x{width}.") + short = min(int(height), int(width)) + for tier in sorted(PIKA_TIERS.values(), key=lambda rule: rule.short_edge): + if short <= tier.short_edge: + return tier.name + return max(PIKA_TIERS.values(), key=lambda rule: rule.short_edge).name + + +def reference_short_edge_for_canvas(height: int, width: int) -> int: + """Reference-image short edge the released tier rule uses for an output canvas.""" + return PIKA_TIERS[tier_for_canvas(height, width)].reference_short_edge + + +def _nearest_aspect(ratio: float, tier: str) -> tuple[int, int, str]: + best: tuple[float, str, tuple[int, int]] | None = None + for aspect in PIKA_ASPECTS: + family_ratio = aspect_ratio(*_parse_aspect(aspect)) + distance = abs(math.log(ratio / family_ratio)) + if best is None or distance < best[0]: + best = (distance, aspect, canvas_for_aspect(aspect, tier)) + # PIKA_ASPECTS is ordered by family; equal distances keep the earlier one. + assert best is not None + height, width = best[2] + return height, width, best[1] + + +def nearest_bucket(height: int, width: int, tier: str) -> tuple[int, int, str]: + """Snap ``height x width`` to the closest aspect family of ``tier``.""" + if height <= 0 or width <= 0: + raise ValueError(f"A canvas must be positive, got {height}x{width}.") + ratio = width / height + if not PIKA_MIN_ASPECT_RATIO - 1e-9 <= ratio <= PIKA_MAX_ASPECT_RATIO + 1e-9: + raise ValueError(f"Pika H3 Flex supports aspect ratios from 1:4 to 4:1, got {width}x{height}.") + return _nearest_aspect(ratio, tier) + + +def resolve_canvas( + height: int | None, + width: int | None, + *, + default_aspect: str | tuple[float, float] = DEFAULT_ASPECT, + default_tier: str | None = DEFAULT_TIER, +) -> tuple[int, int, str, str]: + """Resolve a request canvas to ``(height, width, tier, aspect)``. + + With neither side set, the default tier and aspect are used. With both + sides set, an exact released bucket is honoured; otherwise the tier follows + the requested short edge (``<= 480`` is 480p, everything else 768p) and the + canvas snaps to the nearest aspect family. One-sided requests are rejected. + """ + if (height is None) != (width is None): + raise ValueError("Pika H3 Flex `height` and `width` must be passed together, or neither.") + if height is None: + tier = resolve_tier(default_tier) + aspect_width, aspect_height = _parse_aspect(default_aspect) + bucket_height, bucket_width = canvas_for_aspect((aspect_width, aspect_height), tier) + aspect = PIKA_ASPECTS[0] + for family in PIKA_ASPECTS: + if canvas_for_aspect(family, tier) == (bucket_height, bucket_width): + aspect = family + break + return bucket_height, bucket_width, tier, aspect + + if not isinstance(height, int) or isinstance(height, bool) or not isinstance(width, int) or isinstance(width, bool): + raise TypeError("Pika H3 Flex `height` and `width` must be integers.") + if height <= 0 or width <= 0: + raise ValueError(f"Pika H3 Flex `height` and `width` must be positive, got {height}x{width}.") + + exact = exact_bucket(height, width) + if exact is not None: + tier, aspect = exact + return height, width, tier, aspect + + tier = "p480" if min(height, width) <= PIKA_TIERS["p480"].short_edge else "p768" + bucket_height, bucket_width, aspect = nearest_bucket(height, width, tier) + return bucket_height, bucket_width, tier, aspect + + +def resolve_frames(num_frames: int) -> int: + """Round a request up to Pika's ``17n + 5`` grid inside the 5-15 s range.""" + if not isinstance(num_frames, int) or isinstance(num_frames, bool): + raise TypeError(f"`num_frames` must be an integer, got {type(num_frames).__name__}.") + if num_frames < 1: + raise ValueError(f"`num_frames` must be positive, got {num_frames}.") + aligned = num_frames + while aligned % 17 != 5: + aligned += 1 + if not PIKA_MIN_ALIGNED_FRAMES <= aligned <= PIKA_MAX_ALIGNED_FRAMES: + raise ValueError(f"Pika H3 Flex generates {PIKA_MIN_DURATION:g}-{PIKA_MAX_DURATION:g} seconds at " + f"{PIKA_FPS} fps; aligned num_frames={aligned} is outside the accepted " + f"{PIKA_MIN_ALIGNED_FRAMES}-{PIKA_MAX_ALIGNED_FRAMES} frame range.") + return aligned + + +def duration_seconds(num_frames: int) -> float: + return int(num_frames) / PIKA_FPS + + +__all__ = [ + "DEFAULT_ASPECT", + "DEFAULT_TIER", + "PIKA_ASPECTS", + "PIKA_CANVAS_MULTIPLE", + "PIKA_FPS", + "PIKA_FRAME_GRID", + "PIKA_MAX_ALIGNED_FRAMES", + "PIKA_MAX_DURATION", + "PIKA_MIN_ALIGNED_FRAMES", + "PIKA_MIN_DURATION", + "PIKA_TIERS", + "PikaTier", + "aspect_ratio", + "bucket_table", + "canvas_for_aspect", + "duration_seconds", + "exact_bucket", + "iter_buckets", + "nearest_bucket", + "reference_short_edge_for_canvas", + "resolve_canvas", + "resolve_frames", + "resolve_tier", + "tier_for_canvas", + "tier_rule", +] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py b/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py new file mode 100644 index 000000000..8bfcba982 --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex pipeline: MiniMax-H3 components, Pika adapter and sampler.""" + +from __future__ import annotations + +import math +from typing import Any + +from fastvideo.configs.pipelines.pika_h3_flex import ( + PIKA_AUDIO_SCHEDULER_SHIFT, + PIKA_VIDEO_SCHEDULER_SHIFT, + PikaH3FlexPipelineConfig, +) +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline +from fastvideo.pipelines.basic.pika_h3_flex.adapter import adapter_hf_patterns, load_adapter, task_for_pipeline +from fastvideo.pipelines.basic.pika_h3_flex.stages.denoising import PikaH3FlexDenoisingStage +from fastvideo.pipelines.basic.pika_h3_flex.stages.input_preparation import PikaH3FlexInputPreparationStage + +logger = init_logger(__name__) + + +class PikaH3FlexBasePipeline(MiniMaxH3BasePipeline): + """Pika H3 Flex's adapter, bucket geometry, and six-forward re-noise sampler.""" + + pipeline_config_cls = PikaH3FlexPipelineConfig + input_preparation_stage_cls = PikaH3FlexInputPreparationStage + denoising_stage_cls = PikaH3FlexDenoisingStage + _ref2va_default = False + _pika_adapter: Any = None + + @classmethod + def get_hf_download_allow_patterns(cls) -> list[str]: + patterns = super().get_hf_download_allow_patterns() + # The concrete class picks the adapter, so a T2VA download does not + # fetch the 2.7 GB Ref2VA adapter (and vice versa). + task = "ref2va" if issubclass(cls, PikaH3FlexRefPipeline) else "fl2va" + return [*patterns, *adapter_hf_patterns(task)] + + def _resolve_pika_adapter(self) -> None: + """Point ``lora_path`` at the released adapter before LoRA layers attach. + + ``_load_config`` runs before ``LoRAPipeline`` reads ``lora_path``, so the + automatic adapter behaves exactly like an explicitly passed one, + including the lazy/sequential transformer paths. + """ + task = task_for_pipeline(self._ref2va) + explicit = self.fastvideo_args.lora_path + source = explicit if explicit is not None else self.model_path + revision = None if explicit is not None else self.fastvideo_args.revision + adapter = load_adapter(source, task=task, revision=revision) + source_kind = "explicit" if explicit is not None else "released" + logger.info("Pika H3 Flex: loaded the %s %s adapter from %s", source_kind, task, adapter.path) + self._pika_adapter = adapter + self.fastvideo_args.lora_path = adapter.path + + def _load_config(self, model_path: str) -> dict[str, Any]: + config = super()._load_config(model_path) + self._resolve_pika_adapter() + return config + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + super().initialize_pipeline(fastvideo_args) + expected = {"video": PIKA_VIDEO_SCHEDULER_SHIFT, "audio": PIKA_AUDIO_SCHEDULER_SHIFT} + for module_name, modality in (("scheduler", "video"), ("audio_scheduler", "audio")): + actual = float(self.get_module(module_name).shift) + target = expected[modality] + if not math.isfinite(actual) or abs(actual - target) > 1e-9: + raise ValueError(f"Pika H3 Flex samples with {modality} shift {target:g}, but " + f"{module_name}/scheduler_config.json declares {actual:g}.") + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self._add_stages(fastvideo_args, ref2va=self._ref2va_default) + + +class PikaH3FlexPipeline(PikaH3FlexBasePipeline): + """T2VA/I2VA/L2VA/FL2VA pipeline through the ``fl2va`` adapter.""" + + _ref2va_default = False + + +class PikaH3FlexRefPipeline(PikaH3FlexBasePipeline): + """Ordered-reference pipeline through the ``ref2va`` adapter and ``transformer_ref``.""" + + _extra_config_module_map = {"transformer": "transformer_ref"} + _ref2va_default = True + + +class PikaH3FlexModularPipeline(PikaH3FlexPipeline): + """Public entry matching the released ``modular_model_index.json`` class name.""" + + +class PikaH3FlexRef2VAModularPipeline(PikaH3FlexRefPipeline): + """Public Ref2VA entry matching the released modular class name.""" + + +EntryClass = [PikaH3FlexModularPipeline, PikaH3FlexRef2VAModularPipeline] + +__all__ = [ + "EntryClass", + "PikaH3FlexBasePipeline", + "PikaH3FlexModularPipeline", + "PikaH3FlexPipeline", + "PikaH3FlexRef2VAModularPipeline", + "PikaH3FlexRefPipeline", +] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/presets.py b/fastvideo/pipelines/basic/pika_h3_flex/presets.py new file mode 100644 index 000000000..179ecc8c2 --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/presets.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference presets for Pika H3 Flex joint video/audio generation.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +# The six-forward ladder is the adapter's training contract. Exposing a stage +# override would let a request silently sample the adapter outside its +# distribution, so the denoise stage accepts no overrides. +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Pika's fixed six-forward re-noise ladder", + allowed_overrides=frozenset(), +) + +_SHARED_DEFAULTS = { + "fps": 24, + "guidance_scale": 1.0, + "batch_cfg": False, + "negative_prompt": "", + "num_inference_steps": 6, + "seed": 0, + "height": 768, + "width": 1344, + "num_frames": 124, +} + +PIKA_H3_FLEX_T2VA = InferencePreset( + name="pika_h3_flex_t2va", + version=1, + model_family="pika_h3_flex", + description="Pika H3 Flex text-to-video with synchronized stereo audio at 768p", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +PIKA_H3_FLEX_FL2VA = InferencePreset( + name="pika_h3_flex_fl2va", + version=1, + model_family="pika_h3_flex", + description="Pika H3 Flex first/last-frame-to-video with synchronized stereo audio", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +PIKA_H3_FLEX_REF2VA = InferencePreset( + name="pika_h3_flex_ref2va", + version=1, + model_family="pika_h3_flex", + description="Pika H3 Flex ordered image/video/audio references to joint video and stereo audio", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +ALL_PRESETS = (PIKA_H3_FLEX_T2VA, PIKA_H3_FLEX_FL2VA, PIKA_H3_FLEX_REF2VA) + +__all__ = [ + "ALL_PRESETS", + "PIKA_H3_FLEX_FL2VA", + "PIKA_H3_FLEX_REF2VA", + "PIKA_H3_FLEX_T2VA", +] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/stages/__init__.py b/fastvideo/pipelines/basic/pika_h3_flex/stages/__init__.py new file mode 100644 index 000000000..2e56535f6 --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/stages/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.pika_h3_flex.stages.denoising import PikaH3FlexDenoisingStage +from fastvideo.pipelines.basic.pika_h3_flex.stages.input_preparation import PikaH3FlexInputPreparationStage + +__all__ = ["PikaH3FlexDenoisingStage", "PikaH3FlexInputPreparationStage"] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/stages/denoising.py b/fastvideo/pipelines/basic/pika_h3_flex/stages/denoising.py new file mode 100644 index 000000000..bd648c6b0 --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/stages/denoising.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex's fixed six-forward re-noise denoiser.""" + +from __future__ import annotations + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import MiniMaxH3DenoisingStage +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +def request_seed(batch: ForwardBatch) -> int: + """The request seed, falling back to the generator's initial seed.""" + if batch.seed is not None: + return int(batch.seed) + generator = batch.generator + if isinstance(generator, list): + generator = generator[0] if generator else None + if generator is None: + return 0 + return int(generator.initial_seed()) + + +class PikaH3FlexDenoisingStage(MiniMaxH3DenoisingStage): + """Six full-model forwards over Pika's clean-time ladder. + + Pika's adapter is not a PDD student: every rung is an ordinary transformer + forward. What differs from MiniMax-H3 is the ladder and the re-noise step, + which draws fresh Gaussian noise from a generator seeded ``seed + 1``. + """ + + def _configure_schedules( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + device: torch.device, + ): + steps = list(fastvideo_args.pipeline_config.pika_denoising_steps) + if batch.num_inference_steps != len(steps): + raise ValueError(f"Pika H3 Flex needs {len(steps)} transformer forwards plus the terminal zero, got " + f"num_inference_steps={batch.num_inference_steps}.") + for scheduler in (self.scheduler, self.audio_scheduler): + if not getattr(scheduler, "stochastic_sampling", False): + raise ValueError("Pika H3 Flex needs the released stochastic (re-noise) scheduler; " + f"{type(scheduler).__name__} did not load stochastic_sampling=true.") + self._set_shifted_schedule(steps, len(steps) + 1, device, label="Pika H3 Flex") + return None + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + generator = torch.Generator(device=get_local_torch_device()).manual_seed(request_seed(batch) + 1) + self.scheduler.set_noise_generator(generator) + self.audio_scheduler.set_noise_generator(generator) + return super().forward(batch, fastvideo_args) + + +__all__ = ["PikaH3FlexDenoisingStage", "request_seed"] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/stages/input_preparation.py b/fastvideo/pipelines/basic/pika_h3_flex/stages/input_preparation.py new file mode 100644 index 000000000..6c771458c --- /dev/null +++ b/fastvideo/pipelines/basic/pika_h3_flex/stages/input_preparation.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex input preparation: bucket canvas, cropped keyframes, tiered references.""" + +from __future__ import annotations + +from typing import Any + +import torch +from PIL import Image, ImageOps + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.vision_utils import load_image +from fastvideo.pipelines.basic.minimax_h3.packing import prepare_keyframe_image +from fastvideo.pipelines.basic.minimax_h3.reference import ( + MiniMaxH3PreparedReference, + prepare_reference, + validate_references, +) +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import ( + MINIMAX_H3_KEYFRAME_ANCHORS_KEY, + MINIMAX_H3_KEYFRAMES_KEY, + MiniMaxH3InputPreparationStage, + prepare_common_request, +) +from fastvideo.pipelines.basic.pika_h3_flex.buckets import ( + PIKA_TIERS, + resolve_canvas, + resolve_frames, + resolve_tier, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +logger = init_logger(__name__) + +# Pika's released reference-video floor: two seconds at 24 fps is 48 frames, +# but the adapter's own validator requires at least 22 decoded frames. +PIKA_MIN_REFERENCE_VIDEO_FRAMES = 22 + + +class PikaH3FlexInputPreparationStage(MiniMaxH3InputPreparationStage): + """Prepare T2VA/I2VA/L2VA/FL2VA or Ref2VA inputs on a Pika bucket.""" + + def __init__(self, vae: Any, audio_vae: Any | None = None, *, ref2va: bool = False) -> None: + super().__init__(vae, audio_vae, ref2va=ref2va) + self._tier = "p768" + + def _request_tier(self, fastvideo_args: FastVideoArgs) -> str: + return resolve_tier(getattr(fastvideo_args.pipeline_config, "pika_resolution_tier", None)) + + def _resolve_canvas( + self, + batch: ForwardBatch, + default_aspect: tuple[int, int], + ) -> tuple[int, int, str, str]: + height, width, tier, aspect = resolve_canvas( + batch.height, + batch.width, + default_aspect=default_aspect, + default_tier=self._tier, + ) + if (batch.height, batch.width) != (height, width): + logger.info("Pika H3 Flex snapped canvas %sx%s to %sx%s (%s %s).", batch.height, batch.width, height, width, + tier, aspect) + ratio = int(self.vae.spatial_compression_ratio) + if height % ratio or width % ratio: + raise ValueError(f"Pika H3 Flex canvas {height}x{width} is not divisible by VAE ratio {ratio}.") + return height, width, tier, aspect + + def _prepare_fl2va(self, batch: ForwardBatch) -> None: + if batch.references: + raise ValueError("Pika H3 Flex references belong to the Ref2VA pipeline.") + + if batch.image_path is not None: + batch.pil_image = load_image(batch.image_path) + for name, image in (("pil_image", batch.pil_image), ("last_image", batch.last_image)): + if image is not None and not isinstance(image, Image.Image): + raise TypeError(f"Pika H3 Flex `{name}` must be a PIL image, got {type(image).__name__}.") + + raw_keyframes = [(anchor, ImageOps.exif_transpose(image).convert("RGB")) + for anchor, image in (("first", batch.pil_image), ("last", batch.last_image)) + if image is not None] + default_aspect = raw_keyframes[0][1].size if raw_keyframes else (16, 9) + height, width, _, _ = self._resolve_canvas(batch, default_aspect) + num_frames = resolve_frames(batch.num_frames) + ratio = int(self.vae.spatial_compression_ratio) + self._write_target_geometry(batch, height, width, ratio, num_frames) + + batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = tuple(anchor for anchor, _ in raw_keyframes) + # Pika crops both endpoints. MiniMax-H3 stretches the first frame; using + # that fit here would distort every Pika image on a bucketed canvas. + batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [ + prepare_keyframe_image(image, height, width, stretch=False) for _, image in raw_keyframes + ] + + def _prepare_ref2va(self, batch: ForwardBatch) -> None: + if batch.image_path is not None or batch.pil_image is not None or batch.last_image is not None: + raise ValueError("Pika H3 Flex Ref2VA accepts media through `references`, not FL2VA keyframe fields.") + references = validate_references(list(batch.references or [])) + if self.audio_vae is None: + raise RuntimeError("Pika H3 Flex Ref2VA input preparation has no audio VAE.") + + height, width, tier, _ = self._resolve_canvas(batch, (16, 9)) + num_frames = resolve_frames(batch.num_frames) + target_sample_rate = int(self.audio_vae.sampling_rate) + rule = PIKA_TIERS[tier] + prepared: list[MiniMaxH3PreparedReference] = [] + for reference in references: + item = prepare_reference( + reference, + num_frames, + target_sample_rate, + reference_short_edge=rule.reference_short_edge, + reference_canvas_short_edge=rule.short_edge, + reference_canvas_max_pixels=rule.max_pixels, + ) + if (item.media_type == "video" and item.frames is not None + and item.frames.shape[0] < PIKA_MIN_REFERENCE_VIDEO_FRAMES): + raise ValueError(f"A Pika H3 Flex reference video needs at least " + f"{PIKA_MIN_REFERENCE_VIDEO_FRAMES} decoded frames at 24 fps; got " + f"{item.frames.shape[0]}.") + prepared.append(item) + if not prepared: + raise ValueError("Pika H3 Flex Ref2VA requires at least one non-audio reference.") + + batch.references = prepared + self._write_target_geometry(batch, height, width, int(self.vae.spatial_compression_ratio), num_frames) + batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [] + batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = () + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + prepare_common_request(batch) + self._tier = self._request_tier(fastvideo_args) + if self.ref2va: + self._prepare_ref2va(batch) + else: + self._prepare_fl2va(batch) + return batch + + +__all__ = ["PIKA_MIN_REFERENCE_VIDEO_FRAMES", "PikaH3FlexInputPreparationStage"] diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 414183892..d792c677e 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -42,6 +42,7 @@ from fastvideo.configs.pipelines.flux_2 import ( from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.configs.pipelines.pika_h3_flex import PikaH3FlexPipelineConfig from fastvideo.configs.pipelines.mmaudio import ( MMAudioLarge44kV2PipelineConfig, MMAudioLarge44kV2AConfig, @@ -60,6 +61,7 @@ from fastvideo.api.sampling_param import SamplingParam from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam from fastvideo.api.flux import FluxSamplingParam +from fastvideo.api.pika_h3_flex import PikaH3FlexSamplingParam from fastvideo.fastvideo_args import WorkloadType from fastvideo.logger import init_logger @@ -1152,6 +1154,26 @@ def _register_configs() -> None: _register_wan_configs(WAN_MODEL_DEFINITIONS) + # Pika H3 Flex. It reuses MiniMax-H3's components, but its six-forward + # re-noise sampler and bucket geometry are a separate family contract. + register_configs( + sampling_param_cls=PikaH3FlexSamplingParam, + pipeline_config_cls=PikaH3FlexPipelineConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + hf_model_paths=["Pika-Labs/Pika-H3-Flex"], + model_detectors=[ + lambda path: any(token in path.lower() for token in ( + "pika-h3-flex", + "pika_h3_flex", + "pikah3flexmodularpipeline", + "pikah3flexref2vamodularpipeline", + )), + ], + model_family="pika_h3_flex", + default_preset="pika_h3_flex_t2va", + pipeline_cls_name="PikaH3FlexModularPipeline", + ) + # MiniMax H3 register_configs( sampling_param_cls=None, @@ -1350,6 +1372,8 @@ def _register_presets() -> None: ALL_PRESETS as MATRIXGAME3_PRESETS, ) from fastvideo.pipelines.basic.minimax_h3.presets import ( ALL_PRESETS as MINIMAX_H3_PRESETS, ) + from fastvideo.pipelines.basic.pika_h3_flex.presets import ( + ALL_PRESETS as PIKA_H3_FLEX_PRESETS, ) from fastvideo.pipelines.basic.mmaudio.presets import ( ALL_PRESETS as MMAUDIO_PRESETS, ) from fastvideo.pipelines.basic.sd35.presets import ( @@ -1383,6 +1407,7 @@ def _register_presets() -> None: MATRIXGAME2_PRESETS, MATRIXGAME3_PRESETS, MINIMAX_H3_PRESETS, + PIKA_H3_FLEX_PRESETS, MMAUDIO_PRESETS, SD35_PRESETS, STABLE_AUDIO_PRESETS, diff --git a/fastvideo/tests/loader/test_pika_h3_flex_adapter.py b/fastvideo/tests/loader/test_pika_h3_flex_adapter.py new file mode 100644 index 000000000..ec077e065 --- /dev/null +++ b/fastvideo/tests/loader/test_pika_h3_flex_adapter.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex adapter discovery and manifest contract.""" + +import json +from pathlib import Path + +import pytest + +from fastvideo.pipelines.basic.pika_h3_flex import adapter + + +def _manifest(task: str = "fl2va") -> dict: + return { + "format_version": 1, + "model": "pika-flex", + "task": task, + "release_id": "pika-flex-fl2va-v1", + "rank": 128, + "alpha": 128, + "trunk": "transformer_ref" if task == "ref2va" else "transformer", + "sampling": { + "knots": [1000, 750, 500, 250, 100, 40, 0], + "video_shift": 12.0, + "audio_shift": 3.0, + "method": "renoise", + }, + "base": {"repo_id": "MiniMaxAI/MiniMax-H3", "revision": "42ed227ee7df40d41602854ae760620d6eb651fe"}, + "target_modules": ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2"], + "sha256": "0" * 64, + } + + +def _write_release(root: Path, task: str = "fl2va") -> Path: + task_dir = root / task + task_dir.mkdir(parents=True) + (task_dir / "adapter.safetensors").write_bytes(b"") + (task_dir / "manifest.json").write_text(json.dumps(_manifest(task)), encoding="utf-8") + return task_dir + + +def test_manifest_matches_released_contract() -> None: + adapter.validate_adapter_manifest(_manifest(), task="fl2va") + adapter.validate_adapter_manifest(_manifest("ref2va"), task="ref2va") + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("rank", 64, "rank"), + ("alpha", 64, "alpha"), + ("trunk", "transformer_ref", "trunk"), + ("task", "ref2va", "task"), + ("format_version", 2, "format"), + ], +) +def test_manifest_rejects_wrong_release(field: str, value: object, message: str) -> None: + manifest = _manifest() + manifest[field] = value + with pytest.raises(adapter.PikaAdapterError, match=message): + adapter.validate_adapter_manifest(manifest, task="fl2va") + + +def test_manifest_rejects_a_different_sampler() -> None: + manifest = _manifest() + manifest["sampling"]["knots"] = [1000, 500, 0] + with pytest.raises(adapter.PikaAdapterError, match="sampling knots"): + adapter.validate_adapter_manifest(manifest, task="fl2va") + + +def test_load_adapter_finds_the_task_directory(tmp_path: Path) -> None: + task_dir = _write_release(tmp_path) + loaded = adapter.load_adapter(tmp_path, task="fl2va") + assert loaded.path == str(task_dir / "adapter.safetensors") + assert loaded.trunk == "transformer" + assert loaded.manifest["task"] == "fl2va" + + +def test_load_adapter_finds_a_task_directory_or_adapter_file(tmp_path: Path) -> None: + task_dir = _write_release(tmp_path) + assert adapter.load_adapter(task_dir, task="fl2va").path == str(task_dir / "adapter.safetensors") + assert adapter.load_adapter(task_dir / "adapter.safetensors", task="fl2va").path == str( + task_dir / "adapter.safetensors") + + +def test_load_adapter_rejects_a_missing_task(tmp_path: Path) -> None: + _write_release(tmp_path, "fl2va") + with pytest.raises(adapter.PikaAdapterError, match="ref2va"): + adapter.load_adapter(tmp_path, task="ref2va") + + +def test_load_adapter_requires_the_manifest(tmp_path: Path) -> None: + task_dir = _write_release(tmp_path) + (task_dir / "manifest.json").unlink() + with pytest.raises(adapter.PikaAdapterError, match="manifest"): + adapter.load_adapter(tmp_path, task="fl2va") diff --git a/fastvideo/tests/stages/test_pika_h3_flex_geometry.py b/fastvideo/tests/stages/test_pika_h3_flex_geometry.py new file mode 100644 index 000000000..cd6c1cb70 --- /dev/null +++ b/fastvideo/tests/stages/test_pika_h3_flex_geometry.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex bucket geometry contract.""" + +import pytest + +from fastvideo.pipelines.basic.pika_h3_flex import buckets + + +def test_released_bucket_tables_match_pika() -> None: + assert buckets.bucket_table("p480") == ( + (960, 416), + (832, 480), + (704, 480), + (640, 480), + (608, 480), + (480, 480), + (480, 608), + (480, 640), + (480, 704), + (480, 832), + (416, 960), + ) + assert buckets.bucket_table("p768") == ( + (1536, 672), + (1344, 768), + (1152, 768), + (1024, 768), + (960, 768), + (768, 768), + (768, 960), + (768, 1024), + (768, 1152), + (768, 1344), + (672, 1536), + ) + + +def test_unset_canvas_uses_default_tier_and_aspect() -> None: + assert buckets.resolve_canvas(None, None) == (768, 1344, "p768", "16:9") + + +@pytest.mark.parametrize( + ("height", "width", "expected"), + [ + (720, 1280, (768, 1344, "p768", "16:9")), + (1080, 1920, (768, 1344, "p768", "16:9")), + (480, 832, (480, 832, "p480", "16:9")), + (1024, 768, (1024, 768, "p768", "3:4")), + (960, 416, (960, 416, "p480", "9:21")), + ], +) +def test_requested_canvas_snaps_to_bucket(height: int, width: int, expected: tuple[int, int, str, str]) -> None: + assert buckets.resolve_canvas(height, width) == expected + + +def test_one_sided_canvas_is_rejected() -> None: + with pytest.raises(ValueError, match="together"): + buckets.resolve_canvas(720, None) + + +def test_tier_reference_short_edge_follows_the_output_canvas() -> None: + assert buckets.reference_short_edge_for_canvas(480, 832) == 1024 + assert buckets.reference_short_edge_for_canvas(768, 1344) == 2048 + + +def test_resolve_frames_aligns_to_the_17n_plus_5_grid() -> None: + assert buckets.resolve_frames(120) == 124 + assert buckets.resolve_frames(124) == 124 + assert buckets.resolve_frames(345) == 345 + assert buckets.resolve_frames(360) == 362 + with pytest.raises(ValueError, match="Pika H3 Flex generates"): + buckets.resolve_frames(107) + + +def test_unknown_tier_is_rejected() -> None: + with pytest.raises(ValueError, match="released 480p and 768p"): + buckets.resolve_tier("1080p") diff --git a/fastvideo/tests/stages/test_pika_h3_flex_scheduler.py b/fastvideo/tests/stages/test_pika_h3_flex_scheduler.py new file mode 100644 index 000000000..7a83d3c39 --- /dev/null +++ b/fastvideo/tests/stages/test_pika_h3_flex_scheduler.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex's re-noise scheduler step.""" + +import torch + +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + + +def _scheduler() -> MiniMaxH3Scheduler: + scheduler = MiniMaxH3Scheduler(shift=12.0, stochastic_sampling=True, shift_provided_sigmas=True) + scheduler.set_timesteps(sigmas=[1.0, 0.5, 0.0], device="cpu") + scheduler.set_noise_generator(torch.Generator().manual_seed(7)) + return scheduler + + +def test_renoised_step_uses_the_next_sigma_and_fresh_noise() -> None: + scheduler = _scheduler() + sample = torch.zeros(2, 3) + velocity = torch.ones(2, 3) + timestep = scheduler.timesteps[0] + output = scheduler.step(velocity, timestep, sample).prev_sample + + expected_noise = torch.randn(sample.shape, generator=torch.Generator().manual_seed(7), dtype=torch.float32) + torch.testing.assert_close(output, 0.5 * torch.ones_like(sample) + 0.5 * expected_noise) + + +def test_terminal_renoised_step_returns_x0_without_drawing_noise() -> None: + scheduler = _scheduler() + _ = scheduler.step(torch.ones(2, 3), scheduler.timesteps[0], torch.zeros(2, 3)) + state_before = scheduler._noise_generator.get_state() + output = scheduler.step(torch.ones(2, 3), scheduler.timesteps[1], torch.zeros(2, 3)).prev_sample + state_after = scheduler._noise_generator.get_state() + torch.testing.assert_close(output, torch.full((2, 3), 0.5)) + assert torch.equal(state_before, state_after) + + +def test_euler_step_is_unchanged_when_stochastic_sampling_is_off() -> None: + scheduler = MiniMaxH3Scheduler(shift=12.0) + scheduler.set_timesteps(sigmas=[1.0, 0.5, 0.0], device="cpu") + sample = torch.full((2, 3), 2.0) + velocity = torch.full((2, 3), 4.0) + output = scheduler.step(velocity, scheduler.timesteps[0], sample).prev_sample + # x0 = 2 + 1 * 4 = 6; ratio = 0.5, so prev = 0.5 * 2 + 0.5 * 6 = 4. + torch.testing.assert_close(output, torch.full((2, 3), 4.0))