From 7e960ddf78e0bd45025479ae1dcc7608ea96d311 Mon Sep 17 00:00:00 2001 From: gushiqiao <975033167@qq.com> Date: Tue, 4 Aug 2026 02:33:38 +0000 Subject: [PATCH 1/4] feat: add native MiniMax-H3 AV inference support --- configs/minimax_h3/README.md | 119 +++ configs/minimax_h3/minimax_h3_fl2av.json | 25 + configs/minimax_h3/minimax_h3_i2av.json | 25 + configs/minimax_h3/minimax_h3_l2av.json | 25 + configs/minimax_h3/minimax_h3_ref2av.json | 25 + configs/minimax_h3/minimax_h3_t2av.json | 25 + lightx2v/infer.py | 35 +- .../audio_encoders/hf/minimax_h3/__init__.py | 1 + .../audio_encoders/hf/minimax_h3/audio_vae.py | 598 ++++++++++++ .../input_encoders/hf/minimax_h3/__init__.py | 3 + .../input_encoders/hf/minimax_h3/qwen3vl.py | 876 ++++++++++++++++++ .../hf/minimax_h3/qwen3vl_vision.py | 216 +++++ .../models/networks/minimax_h3/__init__.py | 1 + .../networks/minimax_h3/infer/__init__.py | 1 + .../networks/minimax_h3/infer/module_io.py | 21 + .../networks/minimax_h3/infer/post_infer.py | 26 + .../networks/minimax_h3/infer/pre_infer.py | 117 +++ .../minimax_h3/infer/transformer_infer.py | 78 ++ lightx2v/models/networks/minimax_h3/model.py | 103 ++ .../models/networks/minimax_h3/packing.py | 317 +++++++ .../networks/minimax_h3/packing_ref2av.py | 336 +++++++ .../networks/minimax_h3/weights/__init__.py | 5 + .../minimax_h3/weights/post_weights.py | 23 + .../minimax_h3/weights/pre_weights.py | 63 ++ .../minimax_h3/weights/transformer_weights.py | 51 + .../models/runners/minimax_h3/__init__.py | 3 + .../runners/minimax_h3/minimax_h3_runner.py | 457 +++++++++ .../models/schedulers/minimax_h3/__init__.py | 3 + .../models/schedulers/minimax_h3/scheduler.py | 192 ++++ .../video_encoders/hf/minimax_h3/__init__.py | 1 + .../video_encoders/hf/minimax_h3/video_vae.py | 818 ++++++++++++++++ .../video_encoders/hf/minimax_h3/weights.py | 158 ++++ lightx2v/pipeline.py | 3 + lightx2v/utils/input_info.py | 20 + lightx2v/utils/set_config.py | 29 +- outputs/minimax_h3_t2av.mp4 | Bin 0 -> 1307723 bytes scripts/minimax_h3/run_minimax_h3_fl2av.sh | 18 + scripts/minimax_h3/run_minimax_h3_i2av.sh | 17 + scripts/minimax_h3/run_minimax_h3_l2av.sh | 17 + scripts/minimax_h3/run_minimax_h3_ref2av.sh | 20 + scripts/minimax_h3/run_minimax_h3_t2av.sh | 32 + 41 files changed, 4901 insertions(+), 2 deletions(-) create mode 100644 configs/minimax_h3/README.md create mode 100644 configs/minimax_h3/minimax_h3_fl2av.json create mode 100644 configs/minimax_h3/minimax_h3_i2av.json create mode 100644 configs/minimax_h3/minimax_h3_l2av.json create mode 100644 configs/minimax_h3/minimax_h3_ref2av.json create mode 100644 configs/minimax_h3/minimax_h3_t2av.json create mode 100644 lightx2v/models/audio_encoders/hf/minimax_h3/__init__.py create mode 100644 lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py create mode 100644 lightx2v/models/input_encoders/hf/minimax_h3/__init__.py create mode 100644 lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py create mode 100644 lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py create mode 100644 lightx2v/models/networks/minimax_h3/__init__.py create mode 100644 lightx2v/models/networks/minimax_h3/infer/__init__.py create mode 100644 lightx2v/models/networks/minimax_h3/infer/module_io.py create mode 100644 lightx2v/models/networks/minimax_h3/infer/post_infer.py create mode 100644 lightx2v/models/networks/minimax_h3/infer/pre_infer.py create mode 100644 lightx2v/models/networks/minimax_h3/infer/transformer_infer.py create mode 100644 lightx2v/models/networks/minimax_h3/model.py create mode 100644 lightx2v/models/networks/minimax_h3/packing.py create mode 100644 lightx2v/models/networks/minimax_h3/packing_ref2av.py create mode 100644 lightx2v/models/networks/minimax_h3/weights/__init__.py create mode 100644 lightx2v/models/networks/minimax_h3/weights/post_weights.py create mode 100644 lightx2v/models/networks/minimax_h3/weights/pre_weights.py create mode 100644 lightx2v/models/networks/minimax_h3/weights/transformer_weights.py create mode 100644 lightx2v/models/runners/minimax_h3/__init__.py create mode 100644 lightx2v/models/runners/minimax_h3/minimax_h3_runner.py create mode 100644 lightx2v/models/schedulers/minimax_h3/__init__.py create mode 100644 lightx2v/models/schedulers/minimax_h3/scheduler.py create mode 100644 lightx2v/models/video_encoders/hf/minimax_h3/__init__.py create mode 100644 lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py create mode 100644 lightx2v/models/video_encoders/hf/minimax_h3/weights.py create mode 100644 outputs/minimax_h3_t2av.mp4 create mode 100755 scripts/minimax_h3/run_minimax_h3_fl2av.sh create mode 100755 scripts/minimax_h3/run_minimax_h3_i2av.sh create mode 100755 scripts/minimax_h3/run_minimax_h3_l2av.sh create mode 100755 scripts/minimax_h3/run_minimax_h3_ref2av.sh create mode 100755 scripts/minimax_h3/run_minimax_h3_t2av.sh diff --git a/configs/minimax_h3/README.md b/configs/minimax_h3/README.md new file mode 100644 index 000000000..2ae07a3e6 --- /dev/null +++ b/configs/minimax_h3/README.md @@ -0,0 +1,119 @@ +# MiniMax-H3 原生音视频任务 + +该集成支持 `t2av`、`i2av`、`l2av`、`fl2av` 和 `ref2av`。运行时不导入 Diffusers，也不需要转换或重新下载权重；DiT、Qwen3-VL 前 50 层与视觉塔、视频/音频 VAE 和 scheduler 都由 LightX2V 原生执行。Transformers 只用于 tokenizer 与官方 Qwen3-VL 像素 processor。 + +## 模型目录 + +推荐继续使用 `/data/nvme6/gushiqiao/models/MiniMax-H3` 作为 `model_path`。根目录至少需要以下组件： + +```text +MiniMax-H3/ +├── transformer/ +├── transformer_ref/ +├── vae/ +├── audio_vae/ +├── text_encoder/ +├── tokenizer/ +└── processor/ +``` + +`text_encoder`、`tokenizer` 和 `processor` 可以是指向 `FL2VA/` 对应目录的软链接，不需要复制权重。 + +基础四任务读取 `transformer/`；`ref2av` 自动读取结构相同但数值不同的 `transformer_ref/`。 + +原生 loader 保留发布权重中的混合精度：主体权重为 BF16，原有的 FP32 投影、时间嵌入和输出头保持 FP32。运行时要求 `DTYPE=BF16`；不要求设置 `SENSITIVE_LAYER_DTYPE`，随附脚本设置它只是为了显式表达精度策略。 + +## 任务与输入 + +- `t2av`：仅 prompt。 +- `i2av`：`--image_path`，首帧条件。 +- `l2av`：`--last_frame_path`，尾帧条件。 +- `fl2av`：同时传首帧与尾帧。 +- `ref2av`：重复传 `--reference image=/path`、`video=/path` 或 `audio=/path`，顺序会被保留。最多 9 图、3 视频、3 个带音频 reference、总计 12 项；不能只有音频。 + +- 当前支持单 prompt、batch size 1。 +- 固定 24 fps，官方时长范围为 5–15 秒。 +- 帧数必须满足 `17 * n + 5`；默认是 124 帧。不满足时会先向上对齐，再检查是否仍在时长范围内。 +- 高和宽都必须是 32 的倍数；默认分辨率为 768×1344。 +- 模型是 guidance-distilled，不使用 CFG，`negative_prompt` 会被忽略。 +- 保存路径必须以 `.mp4` 结尾；输出为 24 fps H.264 视频和 32 kHz AAC 立体声音频。 +- 默认 attention 后端是 `torch_sdpa`。只有环境已经安装并支持 FlashAttention 3 时，才应把配置中的 `attn_type` 改为 `flash_attn3`。 + +当前实现面向单卡，并要求按组件顺序执行 CPU offload：文本编码器 → DiT → 视频 VAE → 音频 VAE。配置中应保持 `cpu_offload`、`text_encoder_cpu_offload` 和 `vae_cpu_offload` 为 `true`，`offload_granularity` 为 `model`。 + +## CLI + +在 LightX2V 仓库根目录运行： + +```bash +MODEL_PATH=/data/nvme6/gushiqiao/models/MiniMax-H3 \ +PROMPT='A cinematic wide shot of ocean waves at sunset, with synchronized natural ambience.' \ +OUTPUT_PATH=outputs/minimax_h3_t2av.mp4 \ +bash scripts/minimax_h3/run_minimax_h3_t2av.sh +``` + +可通过 `CONFIG_JSON`、`SEED`、`LIGHTX2V_PATH` 覆盖相应默认值。默认配置是 `configs/minimax_h3/minimax_h3_t2av.json`。 + +```bash +IMAGE_PATH=first.png bash scripts/minimax_h3/run_minimax_h3_i2av.sh +LAST_FRAME_PATH=last.png bash scripts/minimax_h3/run_minimax_h3_l2av.sh +IMAGE_PATH=first.png LAST_FRAME_PATH=last.png bash scripts/minimax_h3/run_minimax_h3_fl2av.sh +REFERENCES='image=person.png,video=motion.mp4,audio=voice.wav' \ + bash scripts/minimax_h3/run_minimax_h3_ref2av.sh +``` + +首尾帧任务在未显式传 `target_shape` 时由第一张实际输入图决定 768p 画布。`ref2av` 的参考图使用自身 2048 短边画布，参考视频重采样到 24 fps，参考音频统一为 32 kHz 立体声。 + +## Python pipeline + +使用同一份 `config_json` 即可初始化程序化入口： + +```python +import os + +os.environ.setdefault("DTYPE", "BF16") + +from lightx2v.pipeline import LightX2VPipeline + +pipe = LightX2VPipeline( + task="t2av", + model_cls="minimax_h3", + model_path="/data/nvme6/gushiqiao/models/MiniMax-H3", +) +pipe.create_generator( + config_json="configs/minimax_h3/minimax_h3_t2av.json", +) +result = pipe.generate( + seed=42, + prompt="A cinematic wide shot of ocean waves at sunset, with synchronized natural ambience.", + save_result_path="outputs/minimax_h3_t2av.mp4", +) +``` + +设置 `return_result_tensor=True` 时返回视频、立体声音频和采样率，不写输出文件。 + +Python API 的条件参数与 CLI 对应。每个 pipeline 在初始化时固定权重分区；`ref2av` 应单独初始化，并接收有序列表： + +```python +ref_pipe = LightX2VPipeline( + task="ref2av", + model_cls="minimax_h3", + model_path="/data/nvme6/gushiqiao/models/MiniMax-H3", +) +ref_pipe.create_generator( + config_json="configs/minimax_h3/minimax_h3_ref2av.json", +) +ref_pipe.generate( + prompt="Keep the character and camera motion, with matching ambience.", + references=["image=person.png", "video=motion.mp4", "audio=voice.wav"], + save_result_path="outputs/ref2av.mp4", +) +``` + +## 当前未支持 + +- CFG、negative prompt 条件分支和 CFG parallel +- sequence parallel、tensor parallel 及多卡推理 +- DiT 量化、LoRA 和 feature caching +- block offload、`lazy_load`、`unload_modules` 和 warmup +- 转换后的 checkpoint 或非官方 safetensors 权重格式 diff --git a/configs/minimax_h3/minimax_h3_fl2av.json b/configs/minimax_h3/minimax_h3_fl2av.json new file mode 100644 index 000000000..0632d5f0b --- /dev/null +++ b/configs/minimax_h3/minimax_h3_fl2av.json @@ -0,0 +1,25 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "torch_sdpa", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} diff --git a/configs/minimax_h3/minimax_h3_i2av.json b/configs/minimax_h3/minimax_h3_i2av.json new file mode 100644 index 000000000..0632d5f0b --- /dev/null +++ b/configs/minimax_h3/minimax_h3_i2av.json @@ -0,0 +1,25 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "torch_sdpa", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} diff --git a/configs/minimax_h3/minimax_h3_l2av.json b/configs/minimax_h3/minimax_h3_l2av.json new file mode 100644 index 000000000..0632d5f0b --- /dev/null +++ b/configs/minimax_h3/minimax_h3_l2av.json @@ -0,0 +1,25 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "torch_sdpa", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} diff --git a/configs/minimax_h3/minimax_h3_ref2av.json b/configs/minimax_h3/minimax_h3_ref2av.json new file mode 100644 index 000000000..0632d5f0b --- /dev/null +++ b/configs/minimax_h3/minimax_h3_ref2av.json @@ -0,0 +1,25 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "torch_sdpa", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} diff --git a/configs/minimax_h3/minimax_h3_t2av.json b/configs/minimax_h3/minimax_h3_t2av.json new file mode 100644 index 000000000..0632d5f0b --- /dev/null +++ b/configs/minimax_h3/minimax_h3_t2av.json @@ -0,0 +1,25 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "torch_sdpa", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} diff --git a/lightx2v/infer.py b/lightx2v/infer.py index de60fb54b..01ecdc39e 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -18,6 +18,7 @@ from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 from lightx2v.models.runners.ltx2.ltx2_runner import LTX2ARRunner, LTX2Runner # noqa: F401 +from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 from lightx2v.models.runners.motus.motus_runner import MotusRunner # noqa: F401 from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 @@ -122,6 +123,7 @@ def main(): "flux2_dev", "ltx2", "ltx2_ar", + "minimax_h3", "bagel", "seedvr2", "neopp", @@ -140,7 +142,31 @@ def main(): parser.add_argument( "--task", type=str, - choices=["t2v", "i2v", "t2t", "t2i", "ti2t", "ti2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "i2va", "v2av", "ltx2_s2v", "sr", "recon", "i23d"], + choices=[ + "t2v", + "i2v", + "t2t", + "t2i", + "ti2t", + "ti2i", + "i2i", + "flf2v", + "vace", + "animate", + "s2v", + "rs2v", + "t2av", + "i2av", + "l2av", + "fl2av", + "ref2av", + "i2va", + "v2av", + "ltx2_s2v", + "sr", + "recon", + "i23d", + ], default="t2v", ) parser.add_argument("--support_tasks", type=str, nargs="+", default=[], help="Set supported tasks for the model") @@ -159,6 +185,13 @@ def main(): ) parser.add_argument("--state_path", type=str, default="", help="The path to input robot state file for robot i2v/i2va inference.") parser.add_argument("--last_frame_path", type=str, default="", help="The path to last frame file for first-last-frame-to-video (flf2v) task") + parser.add_argument( + "--reference", + dest="references", + action="append", + default=[], + help="Ordered MiniMax-H3 ref2av input: image=/path, video=/path, audio=/path, or a JSON object. Repeat to preserve order.", + ) parser.add_argument( "--audio_path", type=str, diff --git a/lightx2v/models/audio_encoders/hf/minimax_h3/__init__.py b/lightx2v/models/audio_encoders/hf/minimax_h3/__init__.py new file mode 100644 index 000000000..a6c81c7e9 --- /dev/null +++ b/lightx2v/models/audio_encoders/hf/minimax_h3/__init__.py @@ -0,0 +1 @@ +from .audio_vae import AutoencoderKLMiniMaxH3AudioNative, MiniMaxH3AudioVAE diff --git a/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py b/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py new file mode 100644 index 000000000..7b82dafc2 --- /dev/null +++ b/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py @@ -0,0 +1,598 @@ +# Copyright 2025 The MiniMax authors and The HuggingFace Team. All rights reserved. +# Copyright 2026 The LightX2V Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Native MiniMax-H3 waveform VAE encoder/decoder. + +The released audio VAE is mono. H3 represents stereo as two batch items: +``[2, 32, frames] -> [2, 1, samples] -> [1, 2, samples]``. The public +:meth:`decode` API keeps that contract, including per-channel latent +denormalization and FP32 DAC/BigVGAN execution. +""" + +from __future__ import annotations + +import gc +import json +import math +from contextlib import nullcontext +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils import weight_norm + +from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( + SafetensorsSubsetReport, + load_safetensors_subset, +) +from lightx2v_platform.base.global_var import AI_DEVICE + + +def _empty_device_cache(device: torch.device) -> None: + backend = getattr(torch, device.type, None) + if backend is not None and hasattr(backend, "empty_cache"): + backend.empty_cache() + + +def _component_dir(model_path: str | Path, component: str) -> Path: + model_path = Path(model_path) + nested = model_path / component + if nested.is_dir(): + return nested + if model_path.name == component and model_path.is_dir(): + return model_path + raise FileNotFoundError(f"Cannot find MiniMax-H3 {component!r} below {model_path}") + + +def _wn_conv1d(*args, **kwargs) -> nn.Module: + # The original checkpoint uses the legacy weight_g/weight_v spelling. + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Kaiser-windowed sinc filter, arithmetically identical to the release.""" + + half_size = kernel_size // 2 + attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + if kernel_size % 2 == 0: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + filter_ /= filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class MiniMaxH3AudioSnakeBeta(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.zeros(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + alpha = torch.exp(self.alpha.unsqueeze(0).unsqueeze(-1)) + beta = torch.exp(self.beta.unsqueeze(0).unsqueeze(-1)) + return hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2) + + +class MiniMaxH3AudioSnake1d(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, hidden_states): + return hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2) + + +class MiniMaxH3AudioResidualUnit(nn.Module): + def __init__(self, dim: int, dilation: int) -> None: + super().__init__() + self.block = nn.Sequential( + MiniMaxH3AudioSnake1d(dim), + _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=3 * dilation), + MiniMaxH3AudioSnake1d(dim), + _wn_conv1d(dim, dim, kernel_size=1), + ) + + def forward(self, hidden_states): + residual = self.block(hidden_states) + pad = (hidden_states.shape[-1] - residual.shape[-1]) // 2 + if pad > 0: + hidden_states = hidden_states[..., pad:-pad] + return hidden_states + residual + + +class MiniMaxH3AudioEncoderBlock(nn.Module): + def __init__(self, dim: int, stride: int) -> None: + super().__init__() + self.block = nn.Sequential( + MiniMaxH3AudioResidualUnit(dim // 2, 1), + MiniMaxH3AudioResidualUnit(dim // 2, 3), + MiniMaxH3AudioResidualUnit(dim // 2, 9), + MiniMaxH3AudioSnake1d(dim // 2), + _wn_conv1d(dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)), + ) + + def forward(self, hidden_states): + return self.block(hidden_states) + + +class MiniMaxH3AudioEncoder(nn.Module): + def __init__(self, d_model: int, strides: tuple[int, ...], d_latent: int) -> None: + super().__init__() + blocks = [_wn_conv1d(1, d_model, kernel_size=7, padding=3)] + for stride in strides: + d_model *= 2 + blocks.append(MiniMaxH3AudioEncoderBlock(d_model, stride)) + blocks += [MiniMaxH3AudioSnake1d(d_model), _wn_conv1d(d_model, d_latent, kernel_size=3, padding=1)] + self.block = nn.Sequential(*blocks) + + def forward(self, hidden_states): + return self.block(hidden_states) + + +class MiniMaxH3AudioGeGluMlp(nn.Module): + def __init__(self, in_features: int, hidden_features: int) -> None: + super().__init__() + self.norm = nn.LayerNorm(in_features) + self.act = nn.GELU(approximate="tanh") + self.w0 = nn.Linear(in_features, hidden_features) + self.w1 = nn.Linear(in_features, hidden_features) + self.w2 = nn.Linear(hidden_features, in_features) + + def forward(self, hidden_states): + hidden_states = self.norm(hidden_states) + return self.w2(self.act(self.w0(hidden_states)) * self.w1(hidden_states)) + + +class MiniMaxH3AudioCausalAttention(nn.Module): + def __init__(self, in_dim: int, out_dim: int, num_heads: int) -> None: + super().__init__() + self.out_dim = out_dim + self.num_heads = num_heads + self.head_dim = in_dim // num_heads + self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False) + self.q_bias = nn.Parameter(torch.zeros(in_dim)) + self.v_bias = nn.Parameter(torch.zeros(in_dim)) + self.register_buffer("zero_k_bias", torch.zeros(in_dim)) + self.proj = nn.Linear(out_dim, out_dim) + + def forward(self, hidden_states): + batch, length, _ = hidden_states.shape + qkv = F.linear(hidden_states, self.qkv.weight, torch.cat((self.q_bias, self.zero_k_bias, self.v_bias))) + query, key, value = qkv.reshape(batch, length, 3, self.num_heads, self.head_dim).permute(2, 0, 1, 3, 4).unbind(0) + output = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), dropout_p=0.0, is_causal=True).transpose(1, 2) + output = output.mean(dim=2) + output = F.adaptive_avg_pool1d(output, self.out_dim) + return self.proj(output) + + +class MiniMaxH3AudioAttnProjection(nn.Module): + def __init__(self, in_dim: int, out_dim: int, num_heads: int, mlp_ratio: int = 2) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(in_dim) + self.attn = MiniMaxH3AudioCausalAttention(in_dim, out_dim, num_heads) + self.proj = nn.Linear(in_dim, out_dim) + self.norm3 = nn.LayerNorm(in_dim) + self.norm2 = nn.LayerNorm(out_dim) + self.mlp = MiniMaxH3AudioGeGluMlp(out_dim, out_dim * mlp_ratio) + + def forward(self, hidden_states): + hidden_states = self.proj(self.norm3(hidden_states)) + self.attn(self.norm1(hidden_states)) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class MiniMaxH3AudioLowPassFilter1d(nn.Module): + def __init__(self, cutoff: float, half_width: float, stride: int, kernel_size: int) -> None: + super().__init__() + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_channels = hidden_states.shape[1] + hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d( + hidden_states, + self.filter.expand(num_channels, -1, -1), + stride=self.stride, + groups=num_channels, + ) + + +class MiniMaxH3AudioUpSample1d(nn.Module): + def __init__(self, ratio: int, kernel_size: int) -> None: + super().__init__() + self.ratio = ratio + self.stride = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (kernel_size - self.stride + 1) // 2 + self.register_buffer( + "filter", + kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_channels = hidden_states.shape[1] + hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate") + hidden_states = self.ratio * F.conv_transpose1d( + hidden_states, + self.filter.expand(num_channels, -1, -1), + stride=self.stride, + groups=num_channels, + ) + return hidden_states[..., self.pad_left : -self.pad_right] + + +class MiniMaxH3AudioDownSample1d(nn.Module): + def __init__(self, ratio: int, kernel_size: int) -> None: + super().__init__() + self.lowpass = MiniMaxH3AudioLowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=kernel_size, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lowpass(hidden_states) + + +class MiniMaxH3AudioActivation1d(nn.Module): + def __init__(self, activation: nn.Module, ratio: int = 2, kernel_size: int = 12) -> None: + super().__init__() + self.act = activation + self.upsample = MiniMaxH3AudioUpSample1d(ratio, kernel_size) + self.downsample = MiniMaxH3AudioDownSample1d(ratio, kernel_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.upsample(hidden_states) + hidden_states = self.act(hidden_states) + return self.downsample(hidden_states) + + +class MiniMaxH3AudioAMPBlock(nn.Module): + """BigVGAN anti-aliased multi-periodicity residual block.""" + + def __init__(self, channels: int, kernel_size: int, dilation: tuple[int, ...]) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [ + _wn_conv1d( + channels, + channels, + kernel_size, + dilation=value, + padding=(kernel_size * value - value) // 2, + ) + for value in dilation + ] + ) + self.convs2 = nn.ModuleList( + [ + _wn_conv1d( + channels, + channels, + kernel_size, + dilation=1, + padding=(kernel_size - 1) // 2, + ) + for _ in dilation + ] + ) + self.activations = nn.ModuleList([MiniMaxH3AudioActivation1d(MiniMaxH3AudioSnakeBeta(channels)) for _ in range(2 * len(dilation))]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + acts1, acts2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2): + residual = conv1(act1(hidden_states)) + residual = conv2(act2(residual)) + hidden_states = residual + hidden_states + return hidden_states + + +class MiniMaxH3AudioBigVGANDecoder(nn.Module): + def __init__( + self, + in_channels: int, + upsample_initial_channel: int, + upsample_rates: tuple[int, ...], + upsample_kernel_sizes: tuple[int, ...], + resblock_kernel_sizes: tuple[int, ...], + resblock_dilation_sizes: tuple[tuple[int, ...], ...], + ) -> None: + super().__init__() + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + if self.num_upsamples == 0 or self.num_kernels == 0: + raise ValueError("MiniMax-H3 audio decoder requires at least one upsampler and one residual kernel") + + self.conv_pre = _wn_conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3) + + # Preserve the released ``ups..0`` nesting for direct key parity. + self.ups = nn.ModuleList() + for index, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + nn.ConvTranspose1d( + upsample_initial_channel // (2**index), + upsample_initial_channel // (2 ** (index + 1)), + kernel, + rate, + padding=(kernel - rate) // 2, + ) + ) + ] + ) + ) + + self.resblocks = nn.ModuleList() + channels = upsample_initial_channel + for index in range(self.num_upsamples): + channels = upsample_initial_channel // (2 ** (index + 1)) + for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes): + self.resblocks.append(MiniMaxH3AudioAMPBlock(channels, kernel, tuple(dilation))) + + self.activation_post = MiniMaxH3AudioActivation1d(MiniMaxH3AudioSnakeBeta(channels)) + self.conv_post = _wn_conv1d(channels, 1, 7, 1, padding=3, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_pre(hidden_states) + for upsample_index in range(self.num_upsamples): + hidden_states = self.ups[upsample_index][0](hidden_states) + residual = None + for kernel_index in range(self.num_kernels): + block = self.resblocks[upsample_index * self.num_kernels + kernel_index](hidden_states) + residual = block if residual is None else residual + block + hidden_states = residual / self.num_kernels + + hidden_states = self.activation_post(hidden_states) + hidden_states = self.conv_post(hidden_states) + return torch.clamp(hidden_states, min=-1.0, max=1.0) + + +class MiniMaxH3AudioVAE(nn.Module): + """H3 audio VAE that loads encoder and decoder from the original file.""" + + def __init__( + self, + config: dict, + *, + device: str | torch.device | None = None, + cpu_offload: bool = False, + ) -> None: + super().__init__() + self.config = dict(config) + self.execution_device = torch.device(device or AI_DEVICE) + self.cpu_offload = cpu_offload + + encoder_rates = tuple(int(value) for value in config.get("encoder_rates", (2, 4, 4, 5, 5))) + decoder_rates = tuple(int(value) for value in config.get("decoder_rates", (5, 5, 2, 2, 2, 2, 2))) + self.hop_length = math.prod(encoder_rates) + if math.prod(decoder_rates) != self.hop_length: + raise ValueError(f"decoder_rates must upsample by the encoder hop length {self.hop_length}, got {math.prod(decoder_rates)}") + + latent_dim = int(config.get("latent_dim", 2048)) + latent_channels = int(config.get("latent_channels", 32)) + self.sampling_rate = int(config.get("sampling_rate", 32000)) + if latent_dim % latent_channels: + raise ValueError(f"latent_dim ({latent_dim}) must be a multiple of latent_channels ({latent_channels})") + self.encoder = MiniMaxH3AudioEncoder(int(config.get("encoder_dim", 64)), encoder_rates, latent_dim) + self.pre_block = MiniMaxH3AudioAttnProjection(latent_dim, latent_channels, int(config.get("num_attention_heads", 8))) + self.mean_proj = nn.Conv1d(latent_channels, latent_channels, 1) + self.logs_proj = nn.Conv1d(latent_channels, latent_channels, 1) + self.dec_in_proj = nn.Conv1d(latent_channels, latent_dim, 1) + self.decoder = MiniMaxH3AudioBigVGANDecoder( + in_channels=latent_dim, + upsample_initial_channel=int(config.get("decoder_dim", 1024)), + upsample_rates=decoder_rates, + upsample_kernel_sizes=tuple(int(value) for value in config.get("decoder_kernel_sizes", (9, 9, 4, 4, 4, 4, 4))), + resblock_kernel_sizes=tuple(int(value) for value in config.get("resblock_kernel_sizes", (3, 7, 11))), + resblock_dilation_sizes=tuple(tuple(int(value) for value in dilation) for dilation in config.get("resblock_dilation_sizes", ((1, 3, 5), (1, 3, 5), (1, 3, 5)))), + ) + + self.register_buffer("latents_mean", torch.empty(latent_channels), persistent=False) + self.register_buffer("latents_std", torch.empty(latent_channels), persistent=False) + self._reset_runtime_buffers() + self.load_report: SafetensorsSubsetReport | None = None + + def _reset_runtime_buffers(self) -> None: + latent_channels = self.dec_in_proj.in_channels + mean = self.config.get("latents_mean", [0.0] * latent_channels) + std = self.config.get("latents_std", [1.0] * latent_channels) + if len(mean) != latent_channels or len(std) != latent_channels: + raise ValueError(f"Audio latent statistics must contain {latent_channels} values, got mean={len(mean)}, std={len(std)}") + self._buffers["latents_mean"] = torch.tensor(mean, dtype=torch.float32) + self._buffers["latents_std"] = torch.tensor(std, dtype=torch.float32) + + @classmethod + def from_pretrained( + cls, + model_path: str | Path, + *, + device: str | torch.device | None = None, + cpu_offload: bool = False, + ) -> "MiniMaxH3AudioVAE": + vae_dir = _component_dir(model_path, "audio_vae") + with (vae_dir / "config.json").open("r", encoding="utf-8") as handle: + config = json.load(handle) + + with torch.device("meta"): + model = cls(config, device=device, cpu_offload=cpu_offload) + model._reset_runtime_buffers() + model.load_report = load_safetensors_subset(model, vae_dir) + model.eval().requires_grad_(False) + if not cpu_offload: + model.to(model.execution_device) + return model + + def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + mean = self.latents_mean.to(device=latents.device).view(1, -1, 1) + std = self.latents_std.to(device=latents.device).view(1, -1, 1) + return latents.float() * std + mean + + def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + mean = self.latents_mean.to(device=latents.device).view(1, -1, 1) + std = self.latents_std.to(device=latents.device).view(1, -1, 1) + return (latents.float() - mean) / std + + def encode(self, waveform: torch.Tensor, *, return_cpu: bool = True) -> torch.Tensor: + """Encode stereo as two mono batch items and return normalized posterior means.""" + try: + if waveform.ndim == 2: + waveform = waveform.unsqueeze(1) + if waveform.ndim != 3 or waveform.shape[1] != 1: + raise ValueError(f"audio waveform must be [batch,1,samples] or [batch,samples], got {tuple(waveform.shape)}") + device = self._activate() + waveform = waveform.to(device=device, dtype=torch.float32) + right_pad = math.ceil(waveform.shape[-1] / self.hop_length) * self.hop_length - waveform.shape[-1] + if right_pad: + waveform = F.pad(waveform, (0, right_pad)) + with torch.no_grad(): + hidden_states = self.encoder(waveform) + hidden_states = self.pre_block(hidden_states.transpose(1, 2)).transpose(1, 2) + mean = self.mean_proj(hidden_states).float() + # The released Ref2AV path consumes posterior.mode(); logs_proj + # is intentionally not evaluated there. + mean = self.normalize_latents(mean) + return mean.cpu() if return_cpu else mean + finally: + if self.cpu_offload: + self.offload() + + def _activate(self) -> torch.device: + if self.cpu_offload: + self.to(self.execution_device) + device = next(self.parameters()).device + dtype = next(self.parameters()).dtype + if dtype != torch.float32: + raise RuntimeError(f"MiniMax-H3 audio VAE weights must remain float32 for parity; found {dtype}. Move the module by device only, without a dtype cast.") + return device + + def offload(self) -> None: + self.to("cpu") + _empty_device_cache(self.execution_device) + gc.collect() + + def _prepare_stereo_latents( + self, + latents: torch.Tensor, + stereo_batch: bool, + ) -> tuple[torch.Tensor, int | None]: + if latents.ndim == 4: + if not stereo_batch or latents.shape[1] != 2: + raise ValueError(f"rank-4 audio latents must use stereo_batch=True and have shape [batch, 2, channels, frames], got {tuple(latents.shape)}") + batch_size, _, channels, frames = latents.shape + return latents.reshape(batch_size * 2, channels, frames), batch_size + if latents.ndim != 3: + raise ValueError(f"audio latents must have shape [2 * batch, channels, frames] or [batch, 2, channels, frames], got {tuple(latents.shape)}") + if stereo_batch: + if latents.shape[0] % 2: + raise ValueError(f"stereo audio uses two mono batch items; leading dimension must be even, got {latents.shape[0]}") + return latents, latents.shape[0] // 2 + return latents, None + + def _run_decode( + self, + latents: torch.Tensor, + *, + denormalize: bool, + stereo_batch: bool, + return_cpu: bool | None, + ) -> torch.Tensor: + try: + latents, stereo_groups = self._prepare_stereo_latents(latents, stereo_batch) + if latents.shape[1] != self.dec_in_proj.in_channels: + raise ValueError(f"audio latents must have {self.dec_in_proj.in_channels} channels, got {latents.shape[1]}") + device = self._activate() + latents = latents.to(device=device, dtype=torch.float32) + if denormalize: + latents = self.denormalize_latents(latents) + # Disable an ambient CUDA autocast: the released DAC/BigVGAN weights + # and arithmetic stay FP32 (BF16 decodes are roughly 20 dB quieter). + autocast_context = torch.autocast(device_type="cuda", enabled=False) if device.type == "cuda" else nullcontext() + with torch.no_grad(), autocast_context: + decoded = self.decoder(self.dec_in_proj(latents)).float() + + if stereo_groups is not None: + decoded = decoded.squeeze(1).reshape(stereo_groups, 2, decoded.shape[-1]) + + if return_cpu is None: + return_cpu = self.cpu_offload + if return_cpu: + decoded = decoded.cpu() + return decoded + finally: + if self.cpu_offload: + self.offload() + + def decode_raw( + self, + denormalized_latents: torch.Tensor, + *, + stereo_batch: bool = False, + return_cpu: bool | None = None, + ) -> torch.Tensor: + """Decode VAE-space latents using the low-level mono-as-batch contract. + + By default ``[2, 32, T]`` returns ``[2, 1, T * 800]``, exactly like + the released autoencoder. Set ``stereo_batch=True`` to fold each pair + into pipeline-facing ``[B, 2, samples]`` output. + """ + + return self._run_decode( + denormalized_latents, + denormalize=False, + stereo_batch=stereo_batch, + return_cpu=return_cpu, + ) + + def decode( + self, + latents: torch.Tensor, + *, + stereo_batch: bool = True, + return_cpu: bool | None = None, + ) -> torch.Tensor: + """Decode normalized H3 audio latents to a waveform in ``[-1, 1]``. + + ``[2, 32, T]`` follows the official mono-as-batch convention and + returns ``[1, 2, T * 800]``. Batched stereo input may alternatively be + supplied as ``[B, 2, 32, T]``. + """ + + return self._run_decode( + latents, + denormalize=True, + stereo_batch=stereo_batch, + return_cpu=return_cpu, + ) + + +AutoencoderKLMiniMaxH3AudioNative = MiniMaxH3AudioVAE diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/__init__.py b/lightx2v/models/input_encoders/hf/minimax_h3/__init__.py new file mode 100644 index 000000000..067c31d04 --- /dev/null +++ b/lightx2v/models/input_encoders/hf/minimax_h3/__init__.py @@ -0,0 +1,3 @@ +from .qwen3vl import MiniMaxH3Qwen3VLTextEncoder, MiniMaxH3TextEncoder + +__all__ = ["MiniMaxH3Qwen3VLTextEncoder", "MiniMaxH3TextEncoder"] diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py new file mode 100644 index 000000000..813bad8ba --- /dev/null +++ b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py @@ -0,0 +1,876 @@ +"""Native Qwen3-VL conditioner for MiniMax-H3 AV inference. + +MiniMax-H3 reads ``hidden_states[50]`` from the released Qwen3-VL +conditioner. In Transformers' hidden-state convention that is the embedding +output after decoder layers 0 through 49, before the final RMSNorm. This +module executes exactly that prefix with LightX2V weight and operator classes. +The vision tower is loaded lazily for keyframe/reference requests; the last +fourteen decoder layers, final norm, and LM head are never loaded. + +Only the Hugging Face tokenizer and pixel processor are reused. Model tensors +are streamed directly from the official sharded ``text_encoder`` safetensors +checkpoint and do not need an offline conversion. +""" + +import gc +import itertools +import json +import os +from collections import defaultdict +from contextlib import suppress +from pathlib import Path + +import torch +import torch.nn.functional as F +from loguru import logger +from safetensors import safe_open + +from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.common.ops.attn.template import AttnWeightTemplate + +# Import the concrete implementations so their registry decorators run even +# when this encoder is imported outside LightX2V's usual top-level entrypoint. +from lightx2v.common.ops.attn.torch_sdpa import TorchSDPAWeight as _TorchSDPAWeight # noqa: F401 +from lightx2v.common.ops.embedding.embedding_weight import EmbeddingWeight as _EmbeddingWeight # noqa: F401 +from lightx2v.common.ops.mm.mm_weight import MMWeight as _MMWeight # noqa: F401 +from lightx2v.common.ops.norm.rms_norm_weight import RMSWeightFP32Qwen as _RMSWeightFP32Qwen # noqa: F401 +from lightx2v.models.input_encoders.hf.minimax_h3.qwen3vl_vision import MiniMaxH3Qwen3VLVisionTower +from lightx2v.models.networks.minimax_h3.packing import VIDEO_TAG +from lightx2v.models.networks.minimax_h3.packing_ref2av import ( + build_ref2av_presentation, + sample_reference_video_frames, +) +from lightx2v.utils.envs import GET_DTYPE +from lightx2v.utils.registry_factory import ( + ATTN_WEIGHT_REGISTER, + EMBEDDING_WEIGHT_REGISTER, + MM_WEIGHT_REGISTER, + RMS_WEIGHT_REGISTER, +) +from lightx2v_platform.base.global_var import AI_DEVICE + +try: + from transformers import Qwen2TokenizerFast, Qwen3VLProcessor +except (AttributeError, ImportError) as error: + Qwen2TokenizerFast = None + Qwen3VLProcessor = None + _TOKENIZER_IMPORT_ERROR = error +else: + _TOKENIZER_IMPORT_ERROR = None + + +MINIMAX_H3_TEXT_ENCODER_LAYER = 50 +MINIMAX_H3_TEXT_HIDDEN_SIZE = 5120 +MINIMAX_H3_TEXT_NUM_LAYERS = 64 +MINIMAX_H3_TEXT_TAG = 1 + +_CHECKPOINT_PREFIX = "model.language_model" +_EXPECTED_RELEASE_CONFIG = { + "hidden_size": MINIMAX_H3_TEXT_HIDDEN_SIZE, + "intermediate_size": 25600, + "num_hidden_layers": MINIMAX_H3_TEXT_NUM_LAYERS, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 151936, +} + + +def _empty_device_cache(): + """Release cached accelerator allocations without assuming CUDA.""" + with suppress(Exception): + device_module = getattr(torch, torch.device(AI_DEVICE).type) + device_module.empty_cache() + + +def _rotate_half(x): + midpoint = x.shape[-1] // 2 + return torch.cat((-x[..., midpoint:], x[..., :midpoint]), dim=-1) + + +def _repeat_kv(x, num_groups): + """Repeat ``[tokens, kv_heads, dim]`` in Qwen's KV-head order.""" + if num_groups == 1: + return x + tokens, num_kv_heads, head_dim = x.shape + return x[:, :, None, :].expand(tokens, num_kv_heads, num_groups, head_dim).reshape(tokens, num_kv_heads * num_groups, head_dim) + + +class _MiniMaxH3QwenSDPAWeight(AttnWeightTemplate): + """Qwen SDPA preserving the released model's native grouped-query path.""" + + def __init__(self): + super().__init__(None) + + def apply( + self, + q, + k, + v, + drop_rate=0, + attn_mask=None, + causal=False, + softmax_scale=None, + **kwargs, + ): + unbatched = q.ndim == 3 + if unbatched: + q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(q.dtype) + output = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=drop_rate, + is_causal=causal, + scale=softmax_scale, + enable_gqa=True, + ) + output = output.transpose(1, 2) + batch_size, sequence_length, num_heads, head_dim = output.shape + output = output.reshape(batch_size, sequence_length, num_heads * head_dim) + return output.squeeze(0) if unbatched else output + + +class _Qwen3VLMLPWeights(WeightModule): + def __init__(self, prefix): + super().__init__() + self.add_module("gate_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.gate_proj.weight")) + self.add_module("up_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.up_proj.weight")) + self.add_module("down_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.down_proj.weight")) + + def forward(self, hidden_states): + gate = self.gate_proj.apply(hidden_states) + up = self.up_proj.apply(hidden_states) + return self.down_proj.apply(F.silu(gate) * up) + + +class _Qwen3VLAttentionWeights(WeightModule): + def __init__(self, prefix, text_config, attn_type): + super().__init__() + self.num_heads = int(text_config["num_attention_heads"]) + self.num_key_value_heads = int(text_config["num_key_value_heads"]) + self.head_dim = int(text_config["head_dim"]) + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.softmax_scale = self.head_dim**-0.5 + eps = float(text_config["rms_norm_eps"]) + + self.add_module("q_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.q_proj.weight")) + self.add_module("k_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.k_proj.weight")) + self.add_module("v_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.v_proj.weight")) + self.add_module("o_proj", MM_WEIGHT_REGISTER["Default"](f"{prefix}.o_proj.weight")) + self.add_module( + "q_norm", + RMS_WEIGHT_REGISTER["fp32_variance_qwen"](f"{prefix}.q_norm.weight", eps=eps), + ) + self.add_module( + "k_norm", + RMS_WEIGHT_REGISTER["fp32_variance_qwen"](f"{prefix}.k_norm.weight", eps=eps), + ) + self.native_gqa = attn_type == "torch_sdpa" + if self.native_gqa: + self.add_module("calculate", _MiniMaxH3QwenSDPAWeight()) + else: + self.add_module("calculate", ATTN_WEIGHT_REGISTER[attn_type]()) + + def forward(self, hidden_states, position_embeddings): + sequence_length = hidden_states.shape[0] + query = self.q_proj.apply(hidden_states).view(sequence_length, self.num_heads, self.head_dim) + key = self.k_proj.apply(hidden_states).view(sequence_length, self.num_key_value_heads, self.head_dim) + value = self.v_proj.apply(hidden_states).view(sequence_length, self.num_key_value_heads, self.head_dim) + + query = self.q_norm.apply(query) + key = self.k_norm.apply(key) + + cos, sin = position_embeddings + cos = cos[:, None, :] + sin = sin[:, None, :] + query = query * cos + _rotate_half(query) * sin + key = key * cos + _rotate_half(key) * sin + + # The released Qwen model keeps eight KV heads and uses torch SDPA's + # native GQA specialization. Preserve that CUDA kernel path for close + # numerical parity; other common backends require materialized heads. + if not self.native_gqa: + key = _repeat_kv(key, self.num_key_value_groups) + value = _repeat_kv(value, self.num_key_value_groups) + attention_output = self.calculate.apply( + query, + key, + value, + causal=True, + max_seqlen_q=sequence_length, + max_seqlen_kv=sequence_length, + softmax_scale=self.softmax_scale, + ) + return self.o_proj.apply(attention_output) + + +class _Qwen3VLDecoderLayerWeights(WeightModule): + def __init__(self, layer_index, text_config, attn_type): + super().__init__() + prefix = f"{_CHECKPOINT_PREFIX}.layers.{layer_index}" + eps = float(text_config["rms_norm_eps"]) + self.add_module( + "input_layernorm", + RMS_WEIGHT_REGISTER["fp32_variance_qwen"](f"{prefix}.input_layernorm.weight", eps=eps), + ) + self.add_module( + "post_attention_layernorm", + RMS_WEIGHT_REGISTER["fp32_variance_qwen"](f"{prefix}.post_attention_layernorm.weight", eps=eps), + ) + self.add_module( + "self_attn", + _Qwen3VLAttentionWeights(f"{prefix}.self_attn", text_config, attn_type), + ) + self.add_module("mlp", _Qwen3VLMLPWeights(f"{prefix}.mlp")) + + def forward(self, hidden_states, position_embeddings): + residual = hidden_states + hidden_states = self.input_layernorm.apply(hidden_states) + hidden_states = residual + self.self_attn.forward(hidden_states, position_embeddings) + + residual = hidden_states + hidden_states = self.post_attention_layernorm.apply(hidden_states) + return residual + self.mlp.forward(hidden_states) + + +class _Qwen3VLTextBackboneWeights(WeightModule): + """Unbatched native prefix of Qwen3-VL's language backbone.""" + + def __init__(self, text_config, num_layers=MINIMAX_H3_TEXT_ENCODER_LAYER, attn_type="torch_sdpa"): + super().__init__() + self.text_config = text_config + self.num_layers = int(num_layers) + self.hidden_size = int(text_config["hidden_size"]) + self.head_dim = int(text_config["head_dim"]) + self.rope_theta = float(text_config["rope_theta"]) + self.add_module( + "embed_tokens", + EMBEDDING_WEIGHT_REGISTER["Default"](f"{_CHECKPOINT_PREFIX}.embed_tokens.weight"), + ) + self.add_module( + "layers", + WeightModuleList(_Qwen3VLDecoderLayerWeights(index, text_config, attn_type) for index in range(self.num_layers)), + ) + self._weight_modules = self._collect_weight_modules() + + # Qwen3-VL's text-only positions are equal on the temporal, height, and + # width axes. Its interleaved M-RoPE therefore reduces exactly to this + # standard split-half RoPE frequency vector. + self._inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, self.head_dim, 2, dtype=torch.int64).to(torch.float32) / self.head_dim)) + rope_config = text_config.get("rope_parameters") or text_config.get("rope_scaling") or {} + self.mrope_section = tuple(rope_config.get("mrope_section", (24, 20, 20))) + + def _collect_weight_modules(self): + modules = {self.embed_tokens.weight_name: self.embed_tokens} + for layer in self.layers: + leaves = ( + layer.input_layernorm, + layer.post_attention_layernorm, + layer.self_attn.q_proj, + layer.self_attn.k_proj, + layer.self_attn.v_proj, + layer.self_attn.o_proj, + layer.self_attn.q_norm, + layer.self_attn.k_norm, + layer.mlp.gate_proj, + layer.mlp.up_proj, + layer.mlp.down_proj, + ) + modules.update((leaf.weight_name, leaf) for leaf in leaves) + return modules + + def named_weight_modules(self): + return self._weight_modules.items() + + @property + def device(self): + return self.embed_tokens.weight.device + + @property + def dtype(self): + return self.embed_tokens.weight.dtype + + def _position_embeddings(self, hidden_states, position_ids=None): + if position_ids is None: + position_ids = torch.arange(hidden_states.shape[0], device=hidden_states.device)[None].expand(3, -1) + if position_ids.ndim != 2 or position_ids.shape[0] != 3: + raise ValueError(f"Qwen3-VL position_ids must be [3,tokens], got {tuple(position_ids.shape)}") + # Transformers' ``from_pretrained(dtype=...)`` keeps Qwen's + # non-persistent RoPE buffers in FP32 even when parameters are BF16. + # Generate phases and trig values in FP32, then cast at the same output + # boundary as the released conditioner. + inv_freq = self._inv_freq.to(device=hidden_states.device, dtype=torch.float32) + frequencies = position_ids.to(torch.float32)[..., None] * inv_freq[None, None, :] + frequencies_t = frequencies[0].clone() + for dim, offset in enumerate((1, 2), start=1): + frequencies_t[..., slice(offset, self.mrope_section[dim] * 3, 3)] = frequencies[dim, ..., slice(offset, self.mrope_section[dim] * 3, 3)] + frequencies = frequencies_t + embeddings = torch.cat((frequencies, frequencies), dim=-1) + return embeddings.cos().to(hidden_states.dtype), embeddings.sin().to(hidden_states.dtype) + + def forward(self, input_ids, position_ids=None, vision_mask=None, vision_embeds=None, deepstack_embeds=None): + if input_ids.ndim != 1: + raise ValueError(f"MiniMax-H3's native Qwen3-VL backbone expects unbatched token IDs, got {tuple(input_ids.shape)}") + hidden_states = self.embed_tokens.apply(input_ids) + if vision_embeds is not None: + if vision_mask is None or int(vision_mask.sum()) != vision_embeds.shape[0]: + raise ValueError("Qwen3-VL vision placeholder count does not match vision embeddings") + hidden_states = hidden_states.clone() + hidden_states[vision_mask] = vision_embeds.to(hidden_states.device, hidden_states.dtype) + position_embeddings = self._position_embeddings(hidden_states, position_ids) + for layer_index, layer in enumerate(self.layers): + hidden_states = layer.forward(hidden_states, position_embeddings) + if deepstack_embeds is not None and layer_index < len(deepstack_embeds): + hidden_states = hidden_states.clone() + hidden_states[vision_mask] += deepstack_embeds[layer_index].to(hidden_states.device, hidden_states.dtype) + return hidden_states + + def to_cpu(self, non_blocking=False): + """Drop accelerator copies; checkpoint weights are immutable pinned tensors.""" + for _, module in self.named_weight_modules(): + pin_weight = getattr(module, "pin_weight", None) + if pin_weight is not None: + module.weight = pin_weight + elif getattr(module, "weight", None) is not None: + module.to_cpu(non_blocking=non_blocking) + return self + + def to_cuda(self, non_blocking=False): + super().to_cuda(non_blocking=non_blocking) + return self + + +class MiniMaxH3Qwen3VLTextEncoder: + """Encode one text-only request with the native first 50 Qwen3-VL layers.""" + + def __init__(self, config): + self.config = config + text_encoder_cpu_offload = bool(config.get("text_encoder_cpu_offload", config.get("cpu_offload", False))) + if "qwen3vl_cpu_offload" in config and bool(config["qwen3vl_cpu_offload"]) != text_encoder_cpu_offload: + raise ValueError("qwen3vl_cpu_offload cannot override text_encoder_cpu_offload for MiniMax-H3; the runner schedules the native conditioner through text_encoder_cpu_offload") + self.cpu_offload = text_encoder_cpu_offload + self.local_files_only = config.get("local_files_only", True) + self.text_encoder = None + self.vision_encoder = None + self.tokenizer = None + self.processor = None + if config.get("text_encoder_load_on_init", True): + self.load() + + @staticmethod + def _require_tokenizer(): + if Qwen2TokenizerFast is None: + detail = f" Original import error: {_TOKENIZER_IMPORT_ERROR}" if _TOKENIZER_IMPORT_ERROR else "" + raise ImportError("MiniMax-H3 text encoding requires Transformers' Qwen2TokenizerFast; the Qwen model itself is executed natively by LightX2V." + detail) + + def _component_path(self, config_key, subfolder): + if config_key in self.config: + return self.config[config_key] + model_path = self.config.get("model_path") + if not model_path: + raise ValueError(f"MiniMax-H3 requires `model_path` or `{config_key}` in the config") + return os.path.join(model_path, subfolder) + + @staticmethod + def _read_text_config(text_encoder_path): + config_path = Path(text_encoder_path) / "config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 text encoder config was not found: {config_path}") + with config_path.open("r", encoding="utf-8") as handle: + raw_config = json.load(handle) + text_config = dict(raw_config.get("text_config", raw_config)) + + rope_config = text_config.get("rope_parameters") or text_config.get("rope_scaling") or {} + text_config["rope_theta"] = float(rope_config.get("rope_theta", text_config.get("rope_theta", 500000.0))) + text_config.setdefault("rms_norm_eps", 1e-6) + text_config.setdefault("attention_bias", False) + text_config.setdefault("hidden_act", "silu") + return text_config + + @staticmethod + def _read_model_config(text_encoder_path): + with (Path(text_encoder_path) / "config.json").open("r", encoding="utf-8") as handle: + return json.load(handle) + + @staticmethod + def _validate_text_config(text_config): + mismatches = {key: (text_config.get(key), expected) for key, expected in _EXPECTED_RELEASE_CONFIG.items() if text_config.get(key) != expected} + if mismatches: + details = ", ".join(f"{key}={actual!r} (expected {expected!r})" for key, (actual, expected) in mismatches.items()) + raise ValueError(f"MiniMax-H3 requires the released Qwen3-VL conditioner config; {details}") + if text_config.get("hidden_act") != "silu": + raise ValueError(f"MiniMax-H3's native Qwen3-VL path requires hidden_act='silu', got {text_config.get('hidden_act')!r}") + if bool(text_config.get("attention_bias", False)): + raise ValueError("MiniMax-H3's released Qwen3-VL attention projections do not use bias") + + num_heads = int(text_config["num_attention_heads"]) + num_kv_heads = int(text_config["num_key_value_heads"]) + head_dim = int(text_config["head_dim"]) + if num_heads % num_kv_heads != 0: + raise ValueError(f"Qwen3-VL attention heads ({num_heads}) must be divisible by KV heads ({num_kv_heads})") + if head_dim % 2: + raise ValueError(f"Qwen3-VL head_dim must be even for split-half RoPE, got {head_dim}") + + @staticmethod + def _resolve_attn_type(config): + requested = config.get( + "qwen3vl_attn_type", + config.get("qwen_attn_implementation", config.get("qwen3vl_attn_implementation", "torch_sdpa")), + ) + aliases = { + "sdpa": "torch_sdpa", + "flash_attention_2": "flash_attn2", + "flash_attention_3": "flash_attn3", + } + attn_type = aliases.get(requested, requested) + if attn_type not in ATTN_WEIGHT_REGISTER: + available = ", ".join(sorted(ATTN_WEIGHT_REGISTER.keys())) + raise ValueError(f"Unknown qwen3vl_attn_type={requested!r}; available LightX2V attention operators: {available}") + return attn_type + + @staticmethod + def _expected_weight_shapes(text_config): + hidden_size = int(text_config["hidden_size"]) + intermediate_size = int(text_config["intermediate_size"]) + num_heads = int(text_config["num_attention_heads"]) + num_kv_heads = int(text_config["num_key_value_heads"]) + head_dim = int(text_config["head_dim"]) + q_size = num_heads * head_dim + kv_size = num_kv_heads * head_dim + + shapes = { + f"{_CHECKPOINT_PREFIX}.embed_tokens.weight": ( + int(text_config["vocab_size"]), + hidden_size, + ) + } + for layer_index in range(MINIMAX_H3_TEXT_ENCODER_LAYER): + prefix = f"{_CHECKPOINT_PREFIX}.layers.{layer_index}" + shapes.update( + { + f"{prefix}.input_layernorm.weight": (hidden_size,), + f"{prefix}.post_attention_layernorm.weight": (hidden_size,), + f"{prefix}.self_attn.q_proj.weight": (q_size, hidden_size), + f"{prefix}.self_attn.k_proj.weight": (kv_size, hidden_size), + f"{prefix}.self_attn.v_proj.weight": (kv_size, hidden_size), + f"{prefix}.self_attn.o_proj.weight": (hidden_size, q_size), + f"{prefix}.self_attn.q_norm.weight": (head_dim,), + f"{prefix}.self_attn.k_norm.weight": (head_dim,), + f"{prefix}.mlp.gate_proj.weight": (intermediate_size, hidden_size), + f"{prefix}.mlp.up_proj.weight": (intermediate_size, hidden_size), + f"{prefix}.mlp.down_proj.weight": (hidden_size, intermediate_size), + } + ) + return shapes + + @staticmethod + def _checkpoint_weight_map(text_encoder_path, required_names): + root = Path(text_encoder_path) + index_path = root / "model.safetensors.index.json" + if index_path.is_file(): + with index_path.open("r", encoding="utf-8") as handle: + index = json.load(handle) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}") + return weight_map + + safetensor_paths = sorted(root.glob("*.safetensors")) + if not safetensor_paths: + raise FileNotFoundError(f"No model.safetensors.index.json or safetensors files found under {root}") + required_names = set(required_names) + weight_map = {} + for checkpoint_path in safetensor_paths: + with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint: + for name in required_names.intersection(checkpoint.keys()): + weight_map[name] = checkpoint_path.name + return weight_map + + @classmethod + def _load_native_weights(cls, backbone, text_encoder_path, text_config): + modules = dict(backbone.named_weight_modules()) + expected_shapes = cls._expected_weight_shapes(text_config) + if modules.keys() != expected_shapes.keys(): + missing_native = sorted(expected_shapes.keys() - modules.keys()) + unexpected_native = sorted(modules.keys() - expected_shapes.keys()) + raise RuntimeError(f"Native Qwen3-VL weight declaration disagrees with its shape schema: missing={missing_native}, unexpected={unexpected_native}") + + root = Path(text_encoder_path) + weight_map = cls._checkpoint_weight_map(root, modules) + missing = sorted(modules.keys() - weight_map.keys()) + if missing: + preview = ", ".join(missing[:8]) + raise KeyError(f"MiniMax-H3 text encoder checkpoint is missing {len(missing)} required tensors: {preview}") + + by_shard = defaultdict(list) + for name in modules: + by_shard[weight_map[name]].append(name) + + # Header-only preflight avoids allocating tens of GiB before discovering + # a malformed or incompatible tensor near the end of the checkpoint. + checkpoint_dtypes = set() + for shard_name in sorted(by_shard): + shard_path = root / shard_name + if not shard_path.is_file(): + raise FileNotFoundError(f"Safetensors shard from checkpoint index was not found: {shard_path}") + with safe_open(shard_path, framework="pt", device="cpu") as checkpoint: + shard_keys = set(checkpoint.keys()) + for name in by_shard[shard_name]: + if name not in shard_keys: + raise KeyError(f"Checkpoint index maps {name} to {shard_path}, but the tensor is absent") + tensor_slice = checkpoint.get_slice(name) + actual_shape = tuple(tensor_slice.get_shape()) + if actual_shape != expected_shapes[name]: + raise ValueError(f"Unexpected checkpoint shape for {name}: {actual_shape}, expected {expected_shapes[name]}") + checkpoint_dtypes.add(tensor_slice.get_dtype()) + if len(checkpoint_dtypes) != 1: + raise ValueError(f"MiniMax-H3 Qwen3-VL weights must use one floating dtype, got {sorted(checkpoint_dtypes)}") + + logger.info( + "Loading {} native Qwen3-VL tensors (embedding + layers 0..{}) from {} shards", + len(modules), + MINIMAX_H3_TEXT_ENCODER_LAYER - 1, + len(by_shard), + ) + for shard_index, shard_name in enumerate(sorted(by_shard), start=1): + shard_path = root / shard_name + logger.info( + "Streaming MiniMax-H3 text shard {}/{}: {}", + shard_index, + len(by_shard), + shard_path.name, + ) + with safe_open(shard_path, framework="pt", device="cpu") as checkpoint: + for name in by_shard[shard_name]: + one_tensor = {name: checkpoint.get_tensor(name)} + module = modules[name] + if module is backbone.embed_tokens: + # The released embedding table is about 1.55 GiB. + # EmbeddingWeight's generic CPU path requires one + # monolithic pinned allocation, which some CUDA + # runtimes reject with cudaErrorInvalidValue. Keep the + # immutable host copy pageable; to_cuda()/F.embedding + # remain the same common-op path, only the transfer is + # synchronous when pinning is unavailable. + module.pin_weight = one_tensor.pop(name) + else: + module.load(one_tensor) + if one_tensor: + raise RuntimeError(f"LightX2V weight loader did not consume tensor {name}") + + # CPU-loaded common weights keep their canonical copy in pin_weight. + # Activate those copies so the object is usable before/after offload. + backbone.to_cpu() + return checkpoint_dtypes.pop() + + def load_tokenizer(self): + """Load only the permitted Transformers tokenizer dependency.""" + self._require_tokenizer() + if self.tokenizer is None: + tokenizer_path = self._component_path("tokenizer_path", "tokenizer") + logger.info(f"Loading MiniMax-H3 tokenizer from {tokenizer_path}") + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + tokenizer_path, + local_files_only=self.local_files_only, + ) + return self.tokenizer + + def unload_tokenizer(self): + self.tokenizer = None + + def load_processor(self): + if Qwen3VLProcessor is None: + detail = f" Original import error: {_TOKENIZER_IMPORT_ERROR}" if _TOKENIZER_IMPORT_ERROR else "" + raise ImportError("MiniMax-H3 image/video conditioning requires Transformers' Qwen3VLProcessor for pixel preprocessing." + detail) + if self.processor is None: + processor_path = self._component_path("processor_path", "processor") + self.processor = Qwen3VLProcessor.from_pretrained(processor_path, local_files_only=self.local_files_only) + return self.processor + + def unload_processor(self): + self.processor = None + + def load_text_encoder(self): + """Stream the native embedding and first 50 decoder layers.""" + if self.text_encoder is not None: + return self.text_encoder + + text_encoder_path = self._component_path("text_encoder_path", "text_encoder") + text_config = self._read_text_config(text_encoder_path) + self._validate_text_config(text_config) + attn_type = self._resolve_attn_type(self.config) + logger.info(f"Building native MiniMax-H3 Qwen3-VL prefix from {text_encoder_path} with attention operator {attn_type}") + text_encoder = _Qwen3VLTextBackboneWeights( + text_config, + num_layers=MINIMAX_H3_TEXT_ENCODER_LAYER, + attn_type=attn_type, + ) + self._load_native_weights(text_encoder, text_encoder_path, text_config) + if not self.cpu_offload: + text_encoder.to_cuda() + self.text_encoder = text_encoder + return self.text_encoder + + def load_vision_encoder(self): + if self.vision_encoder is not None: + return self.vision_encoder + text_encoder_path = self._component_path("text_encoder_path", "text_encoder") + model_config = self._read_model_config(text_encoder_path) + vision_config = dict(model_config["vision_config"]) + logger.info(f"Building native MiniMax-H3 Qwen3-VL vision tower from {text_encoder_path}") + self.vision_encoder = MiniMaxH3Qwen3VLVisionTower.from_pretrained(text_encoder_path, vision_config) + return self.vision_encoder + + def unload_text_encoder(self): + text_encoder = self.text_encoder + self.text_encoder = None + if text_encoder is not None: + del text_encoder + gc.collect() + _empty_device_cache() + + def unload_vision_encoder(self): + vision_encoder = self.vision_encoder + self.vision_encoder = None + if vision_encoder is not None: + del vision_encoder + gc.collect() + _empty_device_cache() + + def to_cpu(self): + if self.text_encoder is not None: + self.text_encoder.to_cpu() + _empty_device_cache() + gc.collect() + return self + + def load(self): + self.load_tokenizer() + self.load_text_encoder() + return self + + def unload(self): + self.unload_vision_encoder() + self.unload_text_encoder() + self.unload_processor() + self.unload_tokenizer() + + def _ensure_loaded(self): + if self.tokenizer is None: + self.load_tokenizer() + if self.text_encoder is None: + self.load_text_encoder() + + def _prepare_t2av_input_ids(self, prompt, device): + if not isinstance(prompt, str): + raise TypeError(f"MiniMax-H3 T2AV expects one prompt string, got {type(prompt).__name__}") + + # Match the upstream conditioner: prompt verbatim, no chat template, + # normalization, padding, or tokenizer-added special tokens. + token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"] + if not token_ids: + raise ValueError("MiniMax-H3 T2AV prompt must encode to at least one token") + return torch.tensor(token_ids, dtype=torch.long, device=device) + + @staticmethod + def _get_rope_index(input_ids, mm_token_type_ids, spatial_merge_size, image_grid_thw=None, video_grid_thw=None): + """Unbatched Qwen3-VL M-RoPE index, matching the released conditioner.""" + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0).clone() + video_grid_thw[:, 0] = 1 + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + groups = [] + for key, group in itertools.groupby(enumerate(mm_token_type_ids.tolist()), lambda item: item[1]): + group = list(group) + groups.append((key, group[0][0], group[-1][0] + 1)) + current_position = 0 + positions = [] + for modality, start, end in groups: + if modality == 0: + length = end - start + positions.append(torch.arange(length, device=input_ids.device).view(1, -1).expand(3, -1) + current_position) + current_position += length + continue + grid = next(grid_iters[modality]) + grid_t = int(grid[0]) + grid_h = int(grid[1]) // spatial_merge_size + grid_w = int(grid[2]) // spatial_merge_size + temporal = torch.arange(grid_t, device=input_ids.device) + height = torch.arange(grid_h, device=input_ids.device) + current_position + width = torch.arange(grid_w, device=input_ids.device) + current_position + t_grid, h_grid, w_grid = torch.meshgrid(temporal, height, width, indexing="ij") + block = torch.stack((t_grid, h_grid, w_grid), dim=0).reshape(3, -1) + block[0] += current_position + positions.append(block) + current_position += max(int(grid[1]), int(grid[2])) // spatial_merge_size + result = torch.cat(positions, dim=1) + if result.shape[1] != input_ids.shape[0]: + raise RuntimeError(f"Qwen3-VL M-RoPE produced {result.shape[1]} positions for {input_ids.shape[0]} tokens") + return result + + def _prepare_keyframe_inputs(self, prompt, images): + processor = self.load_processor() + vision = processor.image_processor(images=images, return_tensors="pt") + image_grid_thw = vision["image_grid_thw"] + merge_unit = processor.image_processor.merge_size**2 + token_ids, token_tags = [], [] + for index in range(len(images)): + count = int(image_grid_thw[index].prod()) // merge_unit + label = self.tokenizer(f": ", add_special_tokens=False)["input_ids"] + block = [self.tokenizer.convert_tokens_to_ids("<|vision_start|>")] + block += [self.tokenizer.convert_tokens_to_ids("<|image_pad|>")] * count + block += [self.tokenizer.convert_tokens_to_ids("<|vision_end|>")] + token_ids += label + block + token_tags += [MINIMAX_H3_TEXT_TAG] * len(label) + [VIDEO_TAG] * len(block) + prompt_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"] + token_ids += prompt_ids + token_tags += [MINIMAX_H3_TEXT_TAG] * len(prompt_ids) + return token_ids, token_tags, vision["pixel_values"], image_grid_thw, None, None + + def _prepare_reference_inputs(self, prompt, references): + processor = self.load_processor() + pixel_values = image_grid_thw = None + image_counts = [] + images = [reference.image for reference in references if reference.kind == "image"] + merge_unit = processor.image_processor.merge_size**2 + if images: + vision = processor.image_processor(images=images, return_tensors="pt") + pixel_values, image_grid_thw = vision["pixel_values"], vision["image_grid_thw"] + image_counts = [int(grid.prod()) // merge_unit for grid in image_grid_thw] + pixel_values_videos = video_grid_thw = None + video_counts = [] + videos = [reference for reference in references if reference.kind == "video"] + if videos: + import numpy as np + + sampled = [sample_reference_video_frames(reference.frames) for reference in videos] + for reference, (_, timestamps) in zip(videos, sampled): + reference.block_timestamps = timestamps + vision = processor.video_processor(videos=[np.stack(frames) for frames, _ in sampled], do_sample_frames=False, return_tensors="pt") + pixel_values_videos, video_grid_thw = vision["pixel_values_videos"], vision["video_grid_thw"] + video_counts = [int(grid[1]) * int(grid[2]) // merge_unit for grid in video_grid_thw] + token_ids, token_tags = build_ref2av_presentation(self.tokenizer, prompt, references, image_counts, video_counts) + return token_ids, token_tags, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw + + def _encode_vision(self, input_ids, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw): + if pixel_values is None and pixel_values_videos is None: + return None, None, None + vision_encoder = self.load_vision_encoder().to(AI_DEVICE) + parameter = next(vision_encoder.parameters()) + image_features = image_deepstack = video_features = video_deepstack = None + try: + with torch.no_grad(): + if pixel_values is not None: + image_features, image_deepstack = vision_encoder(pixel_values.to(AI_DEVICE, parameter.dtype), image_grid_thw.to(AI_DEVICE)) + if pixel_values_videos is not None: + video_features, video_deepstack = vision_encoder(pixel_values_videos.to(AI_DEVICE, parameter.dtype), video_grid_thw.to(AI_DEVICE)) + image_token_id = self.tokenizer.convert_tokens_to_ids("<|image_pad|>") + video_token_id = self.tokenizer.convert_tokens_to_ids("<|video_pad|>") + image_mask, video_mask = input_ids == image_token_id, input_ids == video_token_id + vision_mask = image_mask | video_mask + feature_dim = image_features.shape[-1] if image_features is not None else video_features.shape[-1] + combined = torch.empty((int(vision_mask.sum()), feature_dim), device=AI_DEVICE, dtype=parameter.dtype) + image_joint = image_mask[vision_mask] + video_joint = video_mask[vision_mask] + if image_features is not None: + combined[image_joint] = image_features + if video_features is not None: + combined[video_joint] = video_features + deepstack = [] + source = image_deepstack if image_deepstack is not None else video_deepstack + for layer_index in range(len(source)): + one = torch.empty_like(combined) + if image_deepstack is not None: + one[image_joint] = image_deepstack[layer_index] + if video_deepstack is not None: + one[video_joint] = video_deepstack[layer_index] + deepstack.append(one.cpu()) + return vision_mask.cpu(), combined.cpu(), deepstack + finally: + vision_encoder.to("cpu") + _empty_device_cache() + gc.collect() + + @torch.inference_mode() + def infer(self, prompt, image_list=None, references=None): + """Return unbatched ``[tokens, 5120]`` conditioning and text tags.""" + self._ensure_loaded() + try: + # Input encoding happens before DefaultRunner enters its main-model + # try/finally. Keep migration here so a partially failed transfer + # still reaches the conditioner-specific offload cleanup below. + if references is not None: + prepared = self._prepare_reference_inputs(prompt, references) + elif image_list: + prepared = self._prepare_keyframe_inputs(prompt, image_list) + else: + prepared = None + if prepared is None: + input_ids = self._prepare_t2av_input_ids(prompt, "cpu") + token_tags = torch.full((input_ids.shape[0],), MINIMAX_H3_TEXT_TAG, dtype=torch.long) + position_ids = vision_mask = vision_embeds = deepstack = None + else: + token_ids, tags, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw = prepared + input_ids = torch.tensor(token_ids, dtype=torch.long) + token_tags = torch.tensor(tags, dtype=torch.long) + processor = self.load_processor() + mm_types = torch.tensor(processor.create_mm_token_type_ids([token_ids])[0], dtype=torch.long) + position_ids = self._get_rope_index( + input_ids, + mm_types, + processor.image_processor.merge_size, + image_grid_thw, + video_grid_thw, + ) + vision_mask, vision_embeds, deepstack = self._encode_vision(input_ids, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw) + if self.cpu_offload: + self.text_encoder.to_cuda() + device = self.text_encoder.device + input_ids = input_ids.to(device) + prompt_embeds = self.text_encoder.forward( + input_ids, + None if position_ids is None else position_ids.to(device), + None if vision_mask is None else vision_mask.to(device), + None if vision_embeds is None else vision_embeds.to(device), + None if deepstack is None else [value.to(device) for value in deepstack], + ) + expected_shape = (input_ids.shape[0], MINIMAX_H3_TEXT_HIDDEN_SIZE) + if tuple(prompt_embeds.shape) != expected_shape: + raise RuntimeError(f"MiniMax-H3 expected conditioner hidden shape {expected_shape}, but native Qwen3-VL returned {tuple(prompt_embeds.shape)}") + + prompt_embeds = prompt_embeds.to(device=AI_DEVICE, dtype=GET_DTYPE()).contiguous() + return { + "prompt_embeds": prompt_embeds, + "text_token_tags": token_tags.to(prompt_embeds.device), + } + finally: + if self.cpu_offload and self.text_encoder is not None: + try: + self.text_encoder.to_cpu() + except Exception as error: + logger.warning(f"Best-effort MiniMax-H3 text-encoder offload failed: {error}") + _empty_device_cache() + gc.collect() + + +MiniMaxH3TextEncoder = MiniMaxH3Qwen3VLTextEncoder + + +__all__ = [ + "MINIMAX_H3_TEXT_ENCODER_LAYER", + "MINIMAX_H3_TEXT_HIDDEN_SIZE", + "MINIMAX_H3_TEXT_NUM_LAYERS", + "MINIMAX_H3_TEXT_TAG", + "MiniMaxH3Qwen3VLTextEncoder", + "MiniMaxH3TextEncoder", +] diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py new file mode 100644 index 000000000..9362cb098 --- /dev/null +++ b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py @@ -0,0 +1,216 @@ +"""Native Qwen3-VL vision tower used by MiniMax-H3 conditioning.""" + +import json +import math +from collections import defaultdict +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F +from loguru import logger +from safetensors import safe_open + + +def _vision_position_ids(grid_thw, merge_size): + positions = [] + device = grid_thw.device + for t, h, w in grid_thw.tolist(): + hpos, wpos = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij") + shape = (h // merge_size, merge_size, w // merge_size, merge_size) + hpos = hpos.reshape(shape).transpose(1, 2).flatten() + wpos = wpos.reshape(shape).transpose(1, 2).flatten() + positions.append(torch.stack((hpos, wpos), dim=-1).repeat(t, 1)) + return torch.cat(positions) + + +def _vision_cu_seqlens(grid_thw): + values = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(0, dtype=torch.int32) + return F.pad(values, (1, 0), value=0) + + +def _bilinear_indices_weights(grid_thw, side, merge_size): + device = grid_thw.device + index_parts = [[] for _ in range(4)] + weight_parts = [[] for _ in range(4)] + for t, h, w in grid_thw.tolist(): + h_grid = torch.linspace(0, side - 1, h, device=device) + w_grid = torch.linspace(0, side - 1, w, device=device) + h_floor, w_floor = h_grid.int(), w_grid.int() + h_ceil, w_ceil = (h_floor + 1).clamp(max=side - 1), (w_floor + 1).clamp(max=side - 1) + h_frac, w_frac = h_grid - h_floor, w_grid - w_floor + corners = ( + (h_floor[:, None] * side + w_floor[None]).flatten(), + (h_floor[:, None] * side + w_ceil[None]).flatten(), + (h_ceil[:, None] * side + w_floor[None]).flatten(), + (h_ceil[:, None] * side + w_ceil[None]).flatten(), + ) + weights = ( + ((1 - h_frac)[:, None] * (1 - w_frac)[None]).flatten(), + ((1 - h_frac)[:, None] * w_frac[None]).flatten(), + (h_frac[:, None] * (1 - w_frac)[None]).flatten(), + (h_frac[:, None] * w_frac[None]).flatten(), + ) + h_idx = torch.arange(h, device=device).view(h // merge_size, merge_size) + w_idx = torch.arange(w, device=device).view(w // merge_size, merge_size) + reorder = (h_idx[:, :, None, None] * w + w_idx[None, None]).transpose(1, 2).flatten().repeat(t) + for index in range(4): + index_parts[index].append(corners[index][reorder]) + weight_parts[index].append(weights[index][reorder]) + return torch.stack([torch.cat(part) for part in index_parts]), torch.stack([torch.cat(part) for part in weight_parts]) + + +def _rotate_half(value): + first, second = value.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +class _PatchEmbed(nn.Module): + def __init__(self, config): + super().__init__() + self.in_channels = config["in_channels"] + self.temporal_patch_size = config["temporal_patch_size"] + self.patch_size = config["patch_size"] + kernel = (self.temporal_patch_size, self.patch_size, self.patch_size) + self.proj = nn.Conv3d(self.in_channels, config["hidden_size"], kernel_size=kernel, stride=kernel, bias=True) + + def forward(self, pixels): + pixels = pixels.view(-1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size) + return self.proj(pixels.to(self.proj.weight.dtype)).view(-1, self.proj.out_channels) + + +class _VisionAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.num_heads = config["num_heads"] + self.head_dim = config["hidden_size"] // self.num_heads + self.scaling = self.head_dim**-0.5 + self.qkv = nn.Linear(config["hidden_size"], config["hidden_size"] * 3, bias=True) + self.proj = nn.Linear(config["hidden_size"], config["hidden_size"], bias=True) + + def forward(self, hidden_states, cu_seqlens, cos, sin): + length = hidden_states.shape[0] + query, key, value = self.qkv(hidden_states).reshape(length, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0) + query_f, key_f = query.float(), key.float() + cos_f, sin_f = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + query = (query_f * cos_f + _rotate_half(query_f) * sin_f).to(query.dtype) + key = (key_f * cos_f + _rotate_half(key_f) * sin_f).to(key.dtype) + outputs = [] + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): + q = query[start:end].transpose(0, 1).unsqueeze(0) + k = key[start:end].transpose(0, 1).unsqueeze(0) + v = value[start:end].transpose(0, 1).unsqueeze(0) + out = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=0.0, + is_causal=False, + scale=self.scaling, + ) + outputs.append(out.transpose(1, 2).reshape(end - start, -1)) + return self.proj(torch.cat(outputs, dim=0)) + + +class _VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.linear_fc1 = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=True) + self.linear_fc2 = nn.Linear(config["intermediate_size"], config["hidden_size"], bias=True) + + def forward(self, hidden_states): + return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states), approximate="tanh")) + + +class _VisionBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.norm1 = nn.LayerNorm(config["hidden_size"], eps=1e-6) + self.norm2 = nn.LayerNorm(config["hidden_size"], eps=1e-6) + self.attn = _VisionAttention(config) + self.mlp = _VisionMLP(config) + + def forward(self, hidden_states, cu_seqlens, cos, sin): + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), cu_seqlens, cos, sin) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class _PatchMerger(nn.Module): + def __init__(self, config, postshuffle=False): + super().__init__() + merged_size = config["hidden_size"] * config["spatial_merge_size"] ** 2 + self.merged_size = merged_size + self.postshuffle = postshuffle + self.norm = nn.LayerNorm(merged_size if postshuffle else config["hidden_size"], eps=1e-6) + self.linear_fc1 = nn.Linear(merged_size, merged_size) + self.linear_fc2 = nn.Linear(merged_size, config["out_hidden_size"]) + + def forward(self, hidden_states): + if self.postshuffle: + hidden_states = self.norm(hidden_states.view(-1, self.merged_size)) + else: + hidden_states = self.norm(hidden_states).view(-1, self.merged_size) + return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states))) + + +class MiniMaxH3Qwen3VLVisionTower(nn.Module): + def __init__(self, config): + super().__init__() + self.config = dict(config) + self.spatial_merge_size = int(config["spatial_merge_size"]) + self.patch_embed = _PatchEmbed(config) + self.pos_embed = nn.Embedding(config["num_position_embeddings"], config["hidden_size"]) + self.blocks = nn.ModuleList([_VisionBlock(config) for _ in range(config["depth"])]) + self.merger = _PatchMerger(config) + self.deepstack_visual_indexes = list(config["deepstack_visual_indexes"]) + self.deepstack_merger_list = nn.ModuleList([_PatchMerger(config, postshuffle=True) for _ in self.deepstack_visual_indexes]) + head_dim = config["hidden_size"] // config["num_heads"] + self.register_buffer("rotary_inv_freq", 1.0 / (10000.0 ** (torch.arange(0, head_dim // 2, 2).float() / (head_dim // 2))), persistent=False) + + def forward(self, pixels, grid_thw): + grid_thw = grid_thw.to(device=pixels.device) + indices, weights = _bilinear_indices_weights(grid_thw, int(math.sqrt(self.config["num_position_embeddings"])), self.spatial_merge_size) + position_ids = _vision_position_ids(grid_thw, self.spatial_merge_size) + cu_seqlens = _vision_cu_seqlens(grid_thw) + hidden_states = self.patch_embed(pixels) + pos_embed = (self.pos_embed(indices) * weights[:, :, None]).sum(0) + hidden_states = hidden_states + pos_embed.to(hidden_states.dtype) + rotary = (position_ids.unsqueeze(-1) * self.rotary_inv_freq.to(position_ids.device)).flatten(1) + rotary = torch.cat((rotary, rotary), dim=-1) + cos, sin = rotary.cos(), rotary.sin() + deepstack = [] + for layer_index, block in enumerate(self.blocks): + hidden_states = block(hidden_states, cu_seqlens, cos, sin) + if layer_index in self.deepstack_visual_indexes: + merger_index = self.deepstack_visual_indexes.index(layer_index) + deepstack.append(self.deepstack_merger_list[merger_index](hidden_states)) + return self.merger(hidden_states), deepstack + + @classmethod + def from_pretrained(cls, text_encoder_path, vision_config): + root = Path(text_encoder_path) + with torch.device("meta"): + model = cls(vision_config) + with (root / "model.safetensors.index.json").open("r", encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + prefix = "model.visual." + names = {name: shard for name, shard in weight_map.items() if name.startswith(prefix)} + by_shard = defaultdict(list) + for name, shard in names.items(): + by_shard[shard].append(name) + state = {} + for shard, shard_names in by_shard.items(): + logger.info("Loading native Qwen3-VL vision tensors from {}", shard) + with safe_open(root / shard, framework="pt", device="cpu") as checkpoint: + for name in shard_names: + state[name[len(prefix) :]] = checkpoint.get_tensor(name) + missing, unexpected = model.load_state_dict(state, strict=False, assign=True) + # rotary_inv_freq is non-persistent, so every persistent tensor must match. + if missing or unexpected: + raise RuntimeError(f"Qwen3-VL vision checkpoint mismatch: missing={missing}, unexpected={unexpected}") + head_dim = vision_config["hidden_size"] // vision_config["num_heads"] + model.rotary_inv_freq = 1.0 / (10000.0 ** (torch.arange(0, head_dim // 2, 2, dtype=torch.float32) / (head_dim // 2))) + return model.eval().requires_grad_(False) + + +__all__ = ["MiniMaxH3Qwen3VLVisionTower"] diff --git a/lightx2v/models/networks/minimax_h3/__init__.py b/lightx2v/models/networks/minimax_h3/__init__.py new file mode 100644 index 000000000..b54ddbf71 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/__init__.py @@ -0,0 +1 @@ +"""Native MiniMax-H3 network implementation.""" diff --git a/lightx2v/models/networks/minimax_h3/infer/__init__.py b/lightx2v/models/networks/minimax_h3/infer/__init__.py new file mode 100644 index 000000000..88d1d7772 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/__init__.py @@ -0,0 +1 @@ +"""MiniMax-H3 native inference stages.""" diff --git a/lightx2v/models/networks/minimax_h3/infer/module_io.py b/lightx2v/models/networks/minimax_h3/infer/module_io.py new file mode 100644 index 000000000..f34c55e7d --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/module_io.py @@ -0,0 +1,21 @@ +from dataclasses import dataclass + +import torch + + +@dataclass +class MiniMaxH3PreInferOutput: + hidden_states: torch.Tensor + temb: torch.Tensor + timestep_indices: torch.Tensor + adaln_indices: torch.Tensor + rotary_emb: tuple[torch.Tensor, torch.Tensor] + video_indices: torch.Tensor + audio_indices: torch.Tensor + text_indices: torch.Tensor + + +@dataclass +class MiniMaxH3VelocityOutput: + video: torch.Tensor + audio: torch.Tensor diff --git a/lightx2v/models/networks/minimax_h3/infer/post_infer.py b/lightx2v/models/networks/minimax_h3/infer/post_infer.py new file mode 100644 index 000000000..5b88dc639 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/post_infer.py @@ -0,0 +1,26 @@ +import torch.nn.functional as F + +from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3VelocityOutput +from lightx2v.utils.envs import GET_DTYPE + + +class MiniMaxH3PostInfer: + def __init__(self, config): + self.config = config + + def set_scheduler(self, scheduler): + self.scheduler = scheduler + + def infer(self, weights, hidden_states, pre_infer_out): + shift, scale = weights.norm_out_linear.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())).chunk(2, dim=-1) + indices = pre_infer_out.timestep_indices + hidden_states = weights.norm_out.apply(hidden_states) + hidden_states = hidden_states * (1.0 + scale.index_select(0, indices)) + hidden_states = hidden_states + shift.index_select(0, indices) + + # Both released output heads are fp32 and run over all packed rows + # before modality selection. + hidden_states = hidden_states.float() + video = weights.proj_out.apply(hidden_states).index_select(0, pre_infer_out.video_indices) + audio = weights.audio_proj_out.apply(hidden_states).index_select(0, pre_infer_out.audio_indices) + return MiniMaxH3VelocityOutput(video=video, audio=audio) diff --git a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py new file mode 100644 index 000000000..6654169cd --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py @@ -0,0 +1,117 @@ +import math + +import torch +import torch.nn.functional as F + +from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3PreInferOutput +from lightx2v.utils.envs import GET_DTYPE + + +def timestep_embedding(timesteps: torch.Tensor, embedding_dim: int = 256) -> torch.Tensor: + """Diffusers Timesteps(..., flip_sin_to_cos=True, shift=0), reproduced locally.""" + if timesteps.ndim != 1: + raise ValueError(f"timesteps must be one-dimensional, got {tuple(timesteps.shape)}") + half_dim = embedding_dim // 2 + exponent = -math.log(10000) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / half_dim + phases = timesteps[:, None].float() * torch.exp(exponent)[None] + embedding = torch.cat((torch.cos(phases), torch.sin(phases)), dim=-1) + if embedding_dim % 2: + embedding = F.pad(embedding, (0, 1)) + return embedding + + +class MiniMaxH3PreInfer: + def __init__(self, config): + self.config = config + self.num_heads = int(config.get("num_attention_heads", 56)) + self.head_dim = int(config.get("attention_head_dim", 128)) + self.hidden_size = int(config.get("hidden_size", 5376)) + self.rope_freq_dim = int(config.get("rope_freq_dim", 16)) + self.rope_theta = float(config.get("rope_theta", 10000.0)) + self.freq_dim = int(config.get("freq_dim", 256)) + + def set_scheduler(self, scheduler): + self.scheduler = scheduler + + def _attention(self, weights, hidden_states): + q = weights.to_q.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + k = weights.to_k.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + v = weights.to_v.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + q = weights.norm_q.apply(q) + k = weights.norm_k.apply(k) + seq_len = q.shape[0] + cu_seqlens = torch.tensor((0, seq_len), dtype=torch.int32, device=q.device) + out = weights.calculate.apply( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_len, + max_seqlen_kv=seq_len, + causal=False, + ) + return weights.to_out.apply(out.to(GET_DTYPE())) + + @staticmethod + def _ff(weights, hidden_states): + value, gate = weights.in_proj.apply(hidden_states).chunk(2, dim=-1) + return weights.out_proj.apply(value * F.silu(gate)) + + def _refine_text(self, weights, text_embeds): + for block in weights.refiner_blocks: + text_embeds = text_embeds + self._attention(block.attn, block.norm1.apply(text_embeds)) + text_embeds = text_embeds + self._ff(block.ff, block.norm2.apply(text_embeds)) + return weights.refiner_final_norm.apply(text_embeds) + + def _rotary_embedding(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + position_ids = position_ids.to(torch.float32) + inv_freq = 1.0 / ( + self.rope_theta + ** ( + torch.arange( + 0, + 2 * self.rope_freq_dim, + 2, + dtype=torch.float32, + device=position_ids.device, + ) + / (2 * self.rope_freq_dim) + ) + ) + freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + def infer(self, weights, prompt_embeds): + layout = self.scheduler.layout + bulk_dtype = GET_DTYPE() + + video_embeds = weights.proj_in.apply(self.scheduler.video_latents.float()).to(bulk_dtype) + audio_embeds = weights.audio_proj_in.apply(self.scheduler.audio_latents.float()).to(bulk_dtype) + text_embeds = weights.context_embedder.apply(prompt_embeds.to(bulk_dtype)) + text_embeds = self._refine_text(weights, text_embeds) + + hidden_states = text_embeds.new_zeros((layout.sequence_length, self.hidden_size)) + hidden_states.index_copy_(0, layout.text_indices, text_embeds) + hidden_states.index_copy_(0, layout.audio_indices, audio_embeds) + hidden_states.index_copy_(0, layout.video_indices, video_embeds) + + temb = timestep_embedding(self.scheduler.unique_timesteps, self.freq_dim) + temb = weights.time_linear_2.apply(F.silu(weights.time_linear_1.apply(temb.float()))) + timestep_indices = self.scheduler.timestep_indices + adaln_indices = timestep_indices * 3 + layout.token_tags.clamp(min=0) + + return MiniMaxH3PreInferOutput( + hidden_states=hidden_states, + temb=temb, + timestep_indices=timestep_indices, + adaln_indices=adaln_indices, + rotary_emb=self._rotary_embedding(layout.position_ids), + video_indices=layout.video_indices, + audio_indices=layout.audio_indices, + text_indices=layout.text_indices, + ) diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py new file mode 100644 index 000000000..e30f5b78a --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -0,0 +1,78 @@ +import torch +import torch.nn.functional as F + +from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer +from lightx2v.utils.envs import GET_DTYPE + + +def _apply_rotary_emb(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + rotary_dim = cos.shape[-1] + rotary = hidden_states[..., :rotary_dim] + passthrough = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[:, None, :] + sin = sin.to(hidden_states.dtype)[:, None, :] + x1, x2 = rotary.chunk(2, dim=-1) + rotated = torch.cat((-x2, x1), dim=-1) + rotary = rotary * cos + rotated * sin + return torch.cat((rotary, passthrough), dim=-1).contiguous() + + +class MiniMaxH3TransformerInfer(BaseTransformerInfer): + def __init__(self, config): + self.config = config + self.hidden_size = int(config.get("hidden_size", 5376)) + self.num_heads = int(config.get("num_attention_heads", 56)) + self.head_dim = int(config.get("attention_head_dim", 128)) + self.init_compile(config) + + def _attention(self, weights, hidden_states, rotary_emb): + q = weights.to_q.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + k = weights.to_k.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + v = weights.to_v.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + q = _apply_rotary_emb(weights.norm_q.apply(q), *rotary_emb) + k = _apply_rotary_emb(weights.norm_k.apply(k), *rotary_emb) + seq_len = q.shape[0] + cu_seqlens = torch.tensor((0, seq_len), dtype=torch.int32, device=q.device) + out = weights.calculate.apply( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_len, + max_seqlen_kv=seq_len, + causal=False, + ) + return weights.to_out.apply(out.to(GET_DTYPE())) + + @staticmethod + def _ff(weights, hidden_states): + value, gate = weights.in_proj.apply(hidden_states).chunk(2, dim=-1) + return weights.out_proj.apply(value * F.silu(gate)) + + def infer_block(self, weights, hidden_states, pre_infer_out): + # Activation is evaluated in fp32, then cast immediately before the + # checkpoint's bf16 AdaLN projection. + modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())) + modulation = modulation.view(-1, 6 * self.hidden_size) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1) + indices = pre_infer_out.adaln_indices + + residual = hidden_states + normed = weights.norm1.apply(hidden_states) + normed = normed * (1.0 + scale_msa.index_select(0, indices)) + normed = normed + shift_msa.index_select(0, indices) + hidden_states = residual + gate_msa.index_select(0, indices) * self._attention(weights.attn, normed, pre_infer_out.rotary_emb) + + residual = hidden_states + normed = weights.norm2.apply(hidden_states) + normed = normed * (1.0 + scale_mlp.index_select(0, indices)) + normed = normed + shift_mlp.index_select(0, indices) + hidden_states = residual + gate_mlp.index_select(0, indices) * self._ff(weights.ff, normed) + return hidden_states + + def infer(self, block_weights, pre_infer_out): + hidden_states = pre_infer_out.hidden_states + for block_index, block in enumerate(block_weights.blocks): + hidden_states = self.run_block(block_index, block, hidden_states, pre_infer_out) + return hidden_states diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py new file mode 100644 index 000000000..c6772a758 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -0,0 +1,103 @@ +import os + +import torch +from safetensors import safe_open + +from lightx2v.models.networks.base_model import BaseTransformerModel +from lightx2v.models.networks.minimax_h3.infer.post_infer import MiniMaxH3PostInfer +from lightx2v.models.networks.minimax_h3.infer.pre_infer import MiniMaxH3PreInfer +from lightx2v.models.networks.minimax_h3.infer.transformer_infer import MiniMaxH3TransformerInfer +from lightx2v.models.networks.minimax_h3.weights import ( + MiniMaxH3PostWeights, + MiniMaxH3PreWeights, + MiniMaxH3TransformerWeights, +) +from lightx2v.utils.envs import GET_DTYPE + + +class MiniMaxH3Model(BaseTransformerModel): + """LightX2V-native MiniMax-H3 joint audio/video transformer.""" + + pre_weight_class = MiniMaxH3PreWeights + transformer_weight_class = MiniMaxH3TransformerWeights + post_weight_class = MiniMaxH3PostWeights + + def __init__(self, model_path, config, device): + if GET_DTYPE() != torch.bfloat16: + raise ValueError( + "MiniMax-H3 requires DTYPE=BF16. The native loader preserves the released checkpoint's 626 BF16 tensors and 12 FP32 projection/time/head tensors without dtype conversion." + ) + if config.get("seq_parallel", False): + raise NotImplementedError("MiniMax-H3 sequence parallel support is not implemented yet") + if config.get("tensor_parallel", False): + raise NotImplementedError("MiniMax-H3 tensor parallel support is not implemented yet") + if config.get("cfg_parallel", False) or config.get("enable_cfg", False): + raise ValueError("MiniMax-H3 is guidance-distilled and does not have a CFG/unconditional branch") + if config.get("dit_quantized", False): + raise NotImplementedError("MiniMax-H3 currently supports the released mixed BF16/FP32 weights") + if config.get("lora_configs"): + raise NotImplementedError("MiniMax-H3 LoRA loading is not implemented yet") + if config.get("cpu_offload", False) and config.get("offload_granularity", "model") != "model": + raise NotImplementedError("MiniMax-H3 currently supports component/model CPU offload; block offload needs block-sharded weights") + + transformer_path = config.get("dit_original_ckpt") or os.path.join(model_path, "transformer") + super().__init__(transformer_path, config, device) + self.sensitive_layer = { + "proj_in", + "audio_proj_in", + "time_embedder", + "proj_out", + "audio_proj_out", + } + self._init_infer_class() + self._init_weights() + self._init_infer() + + def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer): + """Load the released mixed-precision tensors without generic dtype coercion.""" + del unified_dtype, sensitive_layer + if os.path.splitext(file_path)[-1] != ".safetensors": + raise ValueError(f"MiniMax-H3 native loading expects the released safetensors checkpoint; got {file_path}") + remove_keys = self.remove_keys if hasattr(self, "remove_keys") else [] + preserve_keys = self.preserved_keys if hasattr(self, "preserved_keys") else None + with safe_open(file_path, framework="pt", device=str(self.device)) as source: + return { + key: source.get_tensor(key) + for key in source.keys() + if not any(remove_key in key for remove_key in remove_keys) and (preserve_keys is None or any(preserve_key in key for preserve_key in preserve_keys)) + } + + def _init_infer_class(self): + if self.config.get("feature_caching", "NoCaching") != "NoCaching": + raise NotImplementedError("MiniMax-H3 feature caching is not implemented") + self.pre_infer_class = MiniMaxH3PreInfer + self.transformer_infer_class = MiniMaxH3TransformerInfer + self.post_infer_class = MiniMaxH3PostInfer + + def _init_infer(self): + self.pre_infer = self.pre_infer_class(self.config) + self.transformer_infer = self.transformer_infer_class(self.config) + self.post_infer = self.post_infer_class(self.config) + + @torch.no_grad() + def _infer_cond_uncond(self, inputs, infer_condition=True): + if not infer_condition: + raise ValueError("MiniMax-H3 does not execute an unconditional pass") + prompt_embeds = inputs["text_encoder_output"]["prompt_embeds"] + pre = self.pre_infer.infer(self.pre_weight, prompt_embeds) + hidden_states = self.transformer_infer.infer(self.transformer_weights, pre) + return self.post_infer.infer(self.post_weight, hidden_states, pre) + + @torch.no_grad() + def infer(self, inputs): + output = self._infer_cond_uncond(inputs, infer_condition=True) + self.scheduler.video_noise_pred = output.video + self.scheduler.audio_noise_pred = output.audio + + @torch.no_grad() + def _seq_parallel_pre_process(self, pre_infer_out): + raise NotImplementedError("MiniMax-H3 sequence parallel support is not implemented") + + @torch.no_grad() + def _seq_parallel_post_process(self, output): + raise NotImplementedError("MiniMax-H3 sequence parallel support is not implemented") diff --git a/lightx2v/models/networks/minimax_h3/packing.py b/lightx2v/models/networks/minimax_h3/packing.py new file mode 100644 index 000000000..3490aa27b --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/packing.py @@ -0,0 +1,317 @@ +"""MiniMax-H3 base-task packed-sequence geometry. + +This is a small, runtime-independent port of the released MiniMax-H3 packing +contract. The transformer consumes one unbatched sequence ordered as +``[text | keyframes | stereo audio | video]``. Keeping this code next to the native +network makes the row order, rotary clock, and noise layout explicit. +""" + +from dataclasses import dataclass + +import numpy as np +import torch +from PIL import Image + +VIDEO_TAG = 0 +TEXT_TAG = 1 +AUDIO_TAG = 2 + +FPS = 24 +AUDIO_LATENTS_PER_SECOND = 40 +AUDIO_CHANNELS = 2 +FRAMES_PER_CHUNK = 17 +LATENTS_PER_CHUNK = 5 +CANVAS_MULTIPLE = 32 +SHORT_EDGE = 768 +MAX_PIXELS = 768 * 1344 +MIN_ASPECT_RATIO = 1.0 / 4.0 +MAX_ASPECT_RATIO = 4.0 +MIN_DURATION = 5.0 +MAX_DURATION = 15.0 +PIXEL_MEAN = (0.485, 0.456, 0.406) +PIXEL_STD = (0.229, 0.224, 0.225) +KEYFRAME_NOISE_AUG = 0.999 +KEYFRAME_ENCODE_SEED = 42 + +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32.0 + + +@dataclass(frozen=True) +class MiniMaxH3PackedSequence: + sequence_length: int + position_ids: torch.Tensor + token_tags: torch.Tensor + video_indices: torch.Tensor + audio_indices: torch.Tensor + text_indices: torch.Tensor + num_condition_video_rows: int = 0 + num_condition_audio_rows: int = 0 + + +def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: + """Resolve an aspect ratio to the released 768p/area-capped H3 canvas.""" + if aspect_width <= 0 or aspect_height <= 0: + raise ValueError(f"The aspect ratio must be positive, got {aspect_width}:{aspect_height}") + ratio = aspect_width / aspect_height + if not MIN_ASPECT_RATIO <= ratio <= 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.0: + width, height = SHORT_EDGE * ratio, float(SHORT_EDGE) + else: + width, height = float(SHORT_EDGE), SHORT_EDGE / ratio + area = width * height + if area > MAX_PIXELS: + scale = (MAX_PIXELS / area) ** 0.5 + width, height = width * scale, height * scale + return ( + max(CANVAS_MULTIPLE, round(height / CANVAS_MULTIPLE) * CANVAS_MULTIPLE), + max(CANVAS_MULTIPLE, round(width / CANVAS_MULTIPLE) * CANVAS_MULTIPLE), + ) + + +def prepare_keyframe_image(image: Image.Image, height: int, width: int, stretch: bool) -> Image.Image: + """Stretch the geometry anchor or cover-crop a following keyframe.""" + if image.size == (width, height): + return image + if stretch: + return image.resize((width, height), Image.Resampling.LANCZOS) + scale = max(width / image.size[0], height / image.size[1]) + resized_size = (max(width, round(image.size[0] * scale)), max(height, round(image.size[1] * scale))) + left = max(0, (resized_size[0] - width) // 2) + top = max(0, (resized_size[1] - height) // 2) + return image.resize(resized_size, Image.Resampling.LANCZOS).crop((left, top, left + width, top + height)) + + +def align_num_frames(num_frames: int) -> int: + if num_frames < 1: + raise ValueError(f"target_video_length must be positive, got {num_frames}") + while num_frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK: + num_frames += 1 + return num_frames + + +def video_latent_num_frames(num_frames: int) -> int: + if num_frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK: + raise ValueError(f"MiniMax-H3 frame count must have the form 17*n+5, got {num_frames}") + return (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2 + + +def audio_latent_num_frames(num_frames: int) -> int: + return int(round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND)) + + +def validate_t2av_geometry(num_frames: int, height: int, width: int) -> None: + duration = num_frames / FPS + if not MIN_DURATION <= duration <= MAX_DURATION: + raise ValueError(f"MiniMax-H3 supports {MIN_DURATION:g}-{MAX_DURATION:g}s at {FPS} fps; got {num_frames} frames ({duration:.3f}s)") + if height % CANVAS_MULTIPLE or width % CANVAS_MULTIPLE: + raise ValueError(f"MiniMax-H3 height and width must be multiples of {CANVAS_MULTIPLE}, got {height}x{width}") + video_latent_num_frames(num_frames) + + +def patchify_video_latents(latents: torch.Tensor, patch_size: tuple[int, int, int] = (1, 2, 2)) -> torch.Tensor: + """Convert ``[1,C,F,H,W]`` latents to unbatched frame-major rows.""" + patch_t, patch_h, patch_w = patch_size + batch, channels, frames, height, width = latents.shape + if batch != 1: + raise ValueError(f"T2AV currently supports batch size 1, got {batch}") + if frames % patch_t or height % patch_h or width % patch_w: + raise ValueError(f"latent shape {tuple(latents.shape)} is not divisible by patch {patch_size}") + latents = latents.reshape( + batch, + channels, + frames // patch_t, + patch_t, + height // patch_h, + patch_h, + width // patch_w, + patch_w, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7) + return latents.reshape(-1, channels * patch_t * patch_h * patch_w).contiguous() + + +def unpatchify_video_tokens( + rows: torch.Tensor, + num_latent_frames: int, + latent_height: int, + latent_width: int, + channels: int = 24, + patch_size: tuple[int, int, int] = (1, 2, 2), +) -> torch.Tensor: + patch_t, patch_h, patch_w = patch_size + rows = rows.reshape( + -1, + num_latent_frames // patch_t, + latent_height // patch_h, + latent_width // patch_w, + channels, + patch_t, + patch_h, + patch_w, + ) + rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7) + return rows.reshape(-1, channels, num_latent_frames, latent_height, latent_width).contiguous() + + +def unpack_audio_tokens(rows: torch.Tensor, num_audio_latents: int) -> torch.Tensor: + rows = rows.reshape(AUDIO_CHANNELS, num_audio_latents, rows.shape[-1]) + return rows.permute(0, 2, 1).contiguous() + + +def _spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + # NumPy's endpoint=False arithmetic is part of the published checkpoint + # contract and differs by an ulp from torch.linspace in some shapes. + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def _temporal_position_grid(num_latent_frames: int, origin: float) -> torch.Tensor: + spans = torch.tensor( + [_ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] for i in range(num_latent_frames)], + dtype=torch.float64, + ) + return origin + torch.cat((torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0))) + + +def temporal_position_span(num_latent_frames: int) -> float: + """Exact NumPy-summed rotary span used for last-frame/reference anchors.""" + spans = np.ones(num_latent_frames, dtype=np.float64) * _ROPE_FRAME_RESCALE + for index, multiplier in enumerate(_ROPE_FRAMES_PER_LATENT): + spans[index :: len(_ROPE_FRAMES_PER_LATENT)] *= multiplier + return float(spans.sum()) + + +def build_t2av_packed_sequence( + num_text_tokens: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + patch_size: tuple[int, int, int] = (1, 2, 2), +) -> MiniMaxH3PackedSequence: + """Build the exact padless ``[text | audio | video]`` T2AV layout.""" + return build_packed_sequence( + torch.full((num_text_tokens,), TEXT_TAG, dtype=torch.long), + num_latent_frames, + latent_height, + latent_width, + num_audio_latents, + patch_size, + ) + + +def build_packed_sequence( + text_token_tags: torch.Tensor, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + patch_size: tuple[int, int, int] = (1, 2, 2), + keyframe_anchors: tuple[str, ...] = (), +) -> MiniMaxH3PackedSequence: + """Build ``[Qwen rows | keyframe rows | target audio | target video]``.""" + _, patch_h, patch_w = patch_size + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_text_tokens = int(text_token_tags.shape[0]) + num_condition_rows = len(keyframe_anchors) * rows_per_frame + num_audio_rows = num_audio_latents * AUDIO_CHANNELS + num_video_rows = num_latent_frames * rows_per_frame + condition_start = num_text_tokens + audio_start = condition_start + num_condition_rows + video_start = audio_start + num_audio_rows + sequence_length = video_start + num_video_rows + + position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_position_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_position_grid(latent_width, patch_w, sqrt_area) + frame_grid = torch.stack([grid.reshape(-1) for grid in torch.meshgrid(height_grid, width_grid, indexing="ij")], dim=-1) + + for index, anchor in enumerate(keyframe_anchors): + if anchor == "first": + anchor_time = float(num_text_tokens) + elif anchor == "last": + anchor_time = float(num_text_tokens) + temporal_position_span(num_latent_frames) - _ROPE_FRAME_RESCALE + else: + raise ValueError(f"A keyframe anchor must be 'first' or 'last', got {anchor!r}") + rows = slice(condition_start + index * rows_per_frame, condition_start + (index + 1) * rows_per_frame) + position_ids[rows, 0] = anchor_time + position_ids[rows, 1:] = frame_grid + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(AUDIO_CHANNELS) + position_ids[audio_start:video_start, 2] = torch.cat( + ( + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full((num_audio_latents,), float(width_grid[-1]), dtype=torch.float64), + ) + ) + + video_positions = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video_positions[:, :, 0] = _temporal_position_grid(num_latent_frames, float(num_text_tokens))[:, None] + video_positions[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video_positions.reshape(-1, 3) + + text_indices = torch.arange(num_text_tokens, dtype=torch.long) + audio_indices = torch.arange(audio_start, video_start, dtype=torch.long) + video_indices = torch.cat((torch.arange(condition_start, audio_start, dtype=torch.long), torch.arange(video_start, sequence_length, dtype=torch.long))) + token_tags = torch.empty(sequence_length, dtype=torch.long) + token_tags[text_indices] = text_token_tags.to(torch.long) + token_tags[audio_indices] = AUDIO_TAG + token_tags[video_indices] = VIDEO_TAG + + return MiniMaxH3PackedSequence( + sequence_length=sequence_length, + position_ids=position_ids, + token_tags=token_tags, + video_indices=video_indices, + audio_indices=audio_indices, + text_indices=text_indices, + num_condition_video_rows=num_condition_rows, + num_condition_audio_rows=0, + ) + + +def build_row_timesteps( + layout: MiniMaxH3PackedSequence, + video_timestep: float, + audio_timestep: float, + condition_video_timestep: float | None = None, + condition_audio_timestep: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + row_timesteps = torch.full((layout.sequence_length,), video_timestep, dtype=torch.float32) + if condition_video_timestep is None: + condition_video_timestep = max(video_timestep, KEYFRAME_NOISE_AUG) + row_timesteps[layout.video_indices[: layout.num_condition_video_rows]] = condition_video_timestep + row_timesteps[layout.audio_indices[layout.num_condition_audio_rows :]] = audio_timestep + row_timesteps[layout.audio_indices[: layout.num_condition_audio_rows]] = condition_audio_timestep + return torch.unique(row_timesteps, sorted=True, return_inverse=True) + + +def keyframe_condition_noise( + condition_latent_shapes: tuple[tuple[int, int, int], ...], + patch_size: tuple[int, int, int], + latent_channels: int, + generator: torch.Generator, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Draw one CPU condition-noise tensor per visual condition, in request order.""" + rows = [] + for num_frames, height, width in condition_latent_shapes: + noise = torch.randn( + (1, latent_channels, num_frames, height, width), + generator=generator, + device="cpu", + dtype=dtype, + ) + rows.append(patchify_video_latents(noise, patch_size)) + if not rows: + return torch.empty((0, latent_channels * int(np.prod(patch_size))), dtype=dtype) + return torch.cat(rows) diff --git a/lightx2v/models/networks/minimax_h3/packing_ref2av.py b/lightx2v/models/networks/minimax_h3/packing_ref2av.py new file mode 100644 index 000000000..c068c9dc0 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/packing_ref2av.py @@ -0,0 +1,336 @@ +"""Native MiniMax-H3 omni-reference preprocessing and packed geometry.""" + +import math +import os +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from .packing import ( + AUDIO_CHANNELS, + AUDIO_TAG, + CANVAS_MULTIPLE, + FPS, + FRAMES_PER_CHUNK, + LATENTS_PER_CHUNK, + TEXT_TAG, + VIDEO_TAG, + MiniMaxH3PackedSequence, + _spatial_position_grid, + _temporal_position_grid, + resolve_canvas_size, +) + +REFERENCE_IMAGE_SHORT_EDGE = 2048 +QWEN_VIDEO_SAMPLE_FPS = 2 +QWEN_TEMPORAL_PATCH = 2 +MAX_REFERENCE_IMAGES = 9 +MAX_REFERENCE_VIDEOS = 3 +MAX_REFERENCE_AUDIOS = 3 +MAX_REFERENCES = 12 +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) + + +@dataclass +class MiniMaxH3PreparedReference: + kind: str + has_audio: bool = False + image: Any = None + frames: Any = None + waveform: torch.Tensor | None = None + block_timestamps: list[float] = field(default_factory=list) + num_latent_frames: int = 1 + latent_height: int = 0 + latent_width: int = 0 + num_audio_latents: int = 0 + video_rows: torch.Tensor | None = None + audio_rows: torch.Tensor | None = None + + @property + def num_video_rows(self) -> int: + return self.num_latent_frames * (self.latent_height // 2) * (self.latent_width // 2) + + @property + def num_audio_rows(self) -> int: + return self.num_audio_latents * AUDIO_CHANNELS + + +def _decode_reference_soundtrack(av, container, stream) -> tuple[torch.Tensor, int]: + """Decode one container stream as planar float at its native rate.""" + + sample_rate = int(stream.codec_context.sample_rate) + resampler = av.audio.resampler.AudioResampler( + format="fltp", + layout=stream.layout, + rate=sample_rate, + ) + chunks = [] + for frame in container.decode(stream): + chunks.extend(torch.from_numpy(value.to_ndarray()) for value in resampler.resample(frame)) + chunks.extend(torch.from_numpy(value.to_ndarray()) for value in resampler.resample(None)) + if not chunks: + raise ValueError("The MiniMax-H3 reference audio stream contains no samples") + return torch.cat(chunks, dim=-1).to(torch.float32), sample_rate + + +def decode_reference_video(media) -> tuple[np.ndarray, float, tuple[torch.Tensor, int] | None]: + """Decode local reference video frames and its optional soundtrack.""" + + path = os.fspath(media) + if not os.path.isfile(path): + raise ValueError(f"MiniMax-H3 reference video is not a local file: {path}") + try: + import av + except ImportError as error: + raise ImportError("Decoding a MiniMax-H3 reference video requires PyAV") from error + + with av.open(path) as container: + if not container.streams.video: + raise ValueError(f"No video stream to decode in {path}") + stream = container.streams.video[0] + frames, rotation = [], 0.0 + for frame in container.decode(stream): + rotation = frame.rotation + frames.append(frame.to_ndarray(format="rgb24")) + frame_rate = float(stream.average_rate or stream.guessed_rate) + soundtrack = None + if container.streams.audio: + container.seek(0) + soundtrack = _decode_reference_soundtrack(av, container, container.streams.audio[0]) + + if not frames: + raise ValueError(f"No video frames to decode in {path}") + frames = np.stack(frames) + turns = round(rotation / 90.0) % 4 + if turns: + frames = np.ascontiguousarray(np.rot90(frames, k=-turns, axes=(1, 2))) + return frames, frame_rate, soundtrack + + +def decode_reference_audio(media) -> tuple[torch.Tensor, int]: + """Decode a local audio reference at the sample rate it carries.""" + + path = os.fspath(media) + if not os.path.isfile(path): + raise ValueError(f"MiniMax-H3 reference audio is not a local file: {path}") + try: + import av + except ImportError as error: + raise ImportError("Decoding a MiniMax-H3 reference audio file requires PyAV") from error + + with av.open(path) as container: + if not container.streams.audio: + raise ValueError(f"No audio stream to decode in {path}") + return _decode_reference_soundtrack(av, container, container.streams.audio[0]) + + +def _temporal_position_span(num_latent_frames: int) -> float: + return sum(_ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[index % len(_ROPE_FRAMES_PER_LATENT)] for index in range(num_latent_frames)) + + +def _frame_position_grid(latent_height: int, latent_width: int, patch_h: int, patch_w: int): + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_position_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_position_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + return torch.stack([grid.reshape(-1) for grid in grids], dim=-1), width_grid + + +def _fill_audio_positions(position_ids, rows, num_audio_latents, rotary_time, width_grid): + time = rotary_time + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[rows, 0] = time.repeat(AUDIO_CHANNELS) + position_ids[rows, 2] = torch.cat( + ( + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full((num_audio_latents,), float(width_grid[-1]), dtype=torch.float64), + ) + ) + + +def build_ref2av_packed_sequence( + text_token_tags: torch.Tensor, + references: list[MiniMaxH3PreparedReference], + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + patch_size: tuple[int, int, int] = (1, 2, 2), +) -> MiniMaxH3PackedSequence: + """Build ``[presentation | ordered references | target audio | target video]``.""" + _, patch_h, patch_w = patch_size + num_text_tokens = int(text_token_tags.shape[0]) + num_target_video_rows = num_latent_frames * (latent_height // patch_h) * (latent_width // patch_w) + num_target_audio_rows = num_audio_latents * AUDIO_CHANNELS + num_reference_video_rows = sum(ref.num_video_rows for ref in references if ref.kind != "audio") + num_reference_audio_rows = sum(ref.num_audio_rows for ref in references) + sequence_length = num_text_tokens + num_reference_video_rows + num_reference_audio_rows + num_target_audio_rows + num_target_video_rows + + position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + target_frame_grid, target_width_grid = _frame_position_grid(latent_height, latent_width, patch_h, patch_w) + video_indices, audio_indices = [], [] + cursor = num_text_tokens + rotary_time = float(num_text_tokens) + for reference in references: + if reference.kind == "image": + rows = slice(cursor, cursor + reference.num_video_rows) + cursor = rows.stop + video_indices.append(torch.arange(rows.start, rows.stop)) + frame_grid, _ = _frame_position_grid(reference.latent_height, reference.latent_width, patch_h, patch_w) + position_ids[rows, 0] = rotary_time + position_ids[rows, 1:] = frame_grid + rotary_time += 1.0 + elif reference.kind == "audio": + rows = slice(cursor, cursor + reference.num_audio_rows) + cursor = rows.stop + audio_indices.append(torch.arange(rows.start, rows.stop)) + _fill_audio_positions(position_ids, rows, reference.num_audio_latents, rotary_time, target_width_grid) + rotary_time += float(reference.num_audio_latents) + elif reference.kind == "video": + audio_rows = slice(cursor, cursor + reference.num_audio_rows) + video_rows = slice(audio_rows.stop, audio_rows.stop + reference.num_video_rows) + cursor = video_rows.stop + audio_indices.append(torch.arange(audio_rows.start, audio_rows.stop)) + video_indices.append(torch.arange(video_rows.start, video_rows.stop)) + frame_grid, width_grid = _frame_position_grid(reference.latent_height, reference.latent_width, patch_h, patch_w) + _fill_audio_positions(position_ids, audio_rows, reference.num_audio_latents, rotary_time, width_grid) + frame_time = _temporal_position_grid(reference.num_latent_frames, rotary_time) + position_ids[video_rows, 0] = frame_time.repeat_interleave(frame_grid.shape[0]) + position_ids[video_rows, 1:] = frame_grid.repeat(reference.num_latent_frames, 1) + rotary_time += max(float(reference.num_audio_latents), _temporal_position_span(reference.num_latent_frames)) + else: + raise ValueError(f"Unknown MiniMax-H3 reference kind: {reference.kind!r}") + + audio_start = cursor + video_start = audio_start + num_target_audio_rows + _fill_audio_positions(position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time, target_width_grid) + frame_time = _temporal_position_grid(num_latent_frames, rotary_time) + position_ids[video_start:, 0] = frame_time.repeat_interleave(target_frame_grid.shape[0]) + position_ids[video_start:, 1:] = target_frame_grid.repeat(num_latent_frames, 1) + video_indices = torch.cat(video_indices + [torch.arange(video_start, sequence_length)]) + audio_indices = torch.cat(audio_indices + [torch.arange(audio_start, video_start)]) + text_indices = torch.arange(num_text_tokens) + token_tags = torch.empty(sequence_length, dtype=torch.long) + token_tags[text_indices] = text_token_tags.long() + token_tags[audio_indices] = AUDIO_TAG + token_tags[video_indices] = VIDEO_TAG + return MiniMaxH3PackedSequence( + sequence_length, + position_ids, + token_tags, + video_indices, + audio_indices, + text_indices, + num_reference_video_rows, + num_reference_audio_rows, + ) + + +def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: + if width <= 0 or height <= 0 or width > 4 * height or height > 4 * width: + raise ValueError(f"A reference image must be positive and within 1:4..4:1, got {width}x{height}") + scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + return ( + max(CANVAS_MULTIPLE, round(height * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE), + max(CANVAS_MULTIPLE, round(width * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE), + ) + + +def prepare_reference_image(image: Image.Image, height: int, width: int) -> Image.Image: + return image if image.size == (width, height) else image.resize((width, height), Image.Resampling.LANCZOS) + + +def resample_reference_frames(frames: np.ndarray, fps: float) -> np.ndarray: + if fps <= 0: + raise ValueError(f"A reference video needs positive fps, got {fps}") + if fps == FPS: + return frames + scale = FPS / fps + slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64) + return np.repeat(frames, np.diff(slots, append=math.floor(frames.shape[0] * scale + 0.5)), axis=0) + + +def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray: + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError(f"Reference video must be [F,H,W,3], got {tuple(frames.shape)}") + frames = frames[:num_frames] + height, width = resolve_canvas_size(frames.shape[2], frames.shape[1]) + if frames.shape[1:3] == (height, width): + return frames + return np.stack([np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS)) for frame in frames]) + + +def sample_reference_video_frames(frames: np.ndarray) -> tuple[list[np.ndarray], list[float]]: + stride = FPS / QWEN_VIDEO_SAMPLE_FPS + indices, cursor = [], 0.0 + while round(cursor) < frames.shape[0]: + if not indices or round(cursor) > indices[-1]: + indices.append(round(cursor)) + cursor += stride + timestamps = [index / QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))] + timestamps += [timestamps[-1]] * (-len(timestamps) % QWEN_TEMPORAL_PATCH) + blocks = [(timestamps[index] + timestamps[index + QWEN_TEMPORAL_PATCH - 1]) / 2 for index in range(0, len(timestamps), QWEN_TEMPORAL_PATCH)] + return [frames[index] for index in indices], blocks + + +def prepare_reference_waveform(waveform, sample_rate: int, target_sample_rate: int, max_duration: float): + waveform = torch.as_tensor(waveform) + if waveform.ndim != 2 or waveform.shape[0] not in (1, AUDIO_CHANNELS): + raise ValueError(f"Reference audio must be mono/stereo [C,S], got {tuple(waveform.shape)}") + waveform = waveform.float()[:, : int(max_duration * sample_rate)] + if waveform.shape[0] == 1: + waveform = waveform.expand(AUDIO_CHANNELS, -1).contiguous() + if sample_rate == target_sample_rate: + return waveform + try: + import torchaudio + except ImportError as error: + raise ImportError("Resampling MiniMax-H3 reference audio requires torchaudio") from error + return torchaudio.transforms.Resample(sample_rate, target_sample_rate)(waveform) + + +def trim_reference_num_frames(num_frames: int) -> int: + if num_frames < 1: + raise ValueError("Reference video contains no frames") + return max(1, (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * FRAMES_PER_CHUNK + LATENTS_PER_CHUNK + + +def build_ref2av_presentation(tokenizer, prompt, references, image_token_counts, video_block_token_counts): + token_ids, token_tags = [], [] + + def emit_text(value): + ids = tokenizer(value, add_special_tokens=False)["input_ids"] + token_ids.extend(ids) + token_tags.extend([TEXT_TAG] * len(ids)) + + def emit_vision(pad_token, count): + ids = [tokenizer.convert_tokens_to_ids("<|vision_start|>")] + ids += [tokenizer.convert_tokens_to_ids(pad_token)] * count + ids += [tokenizer.convert_tokens_to_ids("<|vision_end|>")] + token_ids.extend(ids) + token_tags.extend([VIDEO_TAG] * len(ids)) + + counts = {"image": 0, "video": 0, "audio": 0} + for reference in references: + if reference.has_audio: + counts["audio"] += 1 + emit_text(f"