minimax_h3 warmup and compile

configs/minimax_h3/minimax_h3_fl2av_compile.json

@@ -0,0 +1,27 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 124,
+  "target_height": 544,
+  "target_width": 960,
+  "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": "sage_attn2",
+  "rms_type": "sgl-kernel",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "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
+}

configs/minimax_h3/minimax_h3_i2av_compile.json

@@ -0,0 +1,27 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 124,
+  "target_height": 544,
+  "target_width": 960,
+  "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": "sage_attn2",
+  "rms_type": "sgl-kernel",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "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
+}

configs/minimax_h3/minimax_h3_l2av_compile.json

@@ -0,0 +1,27 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 124,
+  "target_height": 544,
+  "target_width": 960,
+  "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": "sage_attn2",
+  "rms_type": "sgl-kernel",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "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
+}

configs/minimax_h3/minimax_h3_ref2av_compile.json

@@ -0,0 +1,27 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 124,
+  "target_height": 544,
+  "target_width": 960,
+  "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": "sage_attn2",
+  "rms_type": "sgl-kernel",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "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
+}

configs/minimax_h3/minimax_h3_t2av_compile.json

@@ -0,0 +1,27 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 124,
+  "target_height": 544,
+  "target_width": 960,
+  "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": "sage_attn2",
+  "rms_type": "sgl-kernel",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "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
+}

lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py

@@ -19,6 +19,13 @@ def __init__(self, config):
         elif offload_granularity != "model":
             raise NotImplementedError(f"MiniMax-H3 does not support offload_granularity={offload_granularity!r}")
 
+    def get_compile_block_key(self, block_idx, block):
+        # block offload
+        if hasattr(self, "offload_manager"):
+            return id(block)
+        # model offload
+        return super().get_compile_block_key(block_idx, block)
+
     def infer_with_blocks_offload(self, blocks, hidden_states, pre_infer_out):
         num_blocks = len(blocks)
         current_stream = torch_device_module.current_stream()

lightx2v/models/networks/minimax_h3/infer/transformer_infer.py

@@ -21,6 +21,7 @@ def __init__(self, config):
             self.tp_rank = 0
         self.num_heads = self.global_num_heads // self.tp_size
         self.head_dim = int(config.get("attention_head_dim", 128))
+        self.infer_dtype = GET_DTYPE()
         if config.get("seq_parallel", False):
             self.seq_p_group = config["device_mesh"].get_group(mesh_dim="seq_p")
             parallel = config.get("parallel", {})
@@ -90,7 +91,7 @@ def _attention(self, weights, hidden_states, pre_infer_out):
             )
             if aux_out is not None:
                 out = torch.cat((aux_out, out), dim=0)
-        return weights.to_out.apply(out.to(GET_DTYPE()))
+        return weights.to_out.apply(out.to(self.infer_dtype))
 
     @staticmethod
     def _ff(weights, hidden_states):
@@ -100,7 +101,7 @@ def _ff(weights, hidden_states):
     def infer_block(self, weights, hidden_states, pre_infer_out):
         # Activation is evaluated in fp32, then cast to the inference dtype
         # immediately before the (possibly quantized) AdaLN projection.
-        modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE()))
+        modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(self.infer_dtype))
         modulation = self._gather_tp_last_dim(modulation)
         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)

lightx2v/models/runners/minimax_h3/minimax_h3_runner.py

@@ -61,6 +61,9 @@ class MiniMaxH3Runner(DefaultRunner):
     graph.
     """
 
+    _WARMUP_RESOLUTIONS = ((480, 480), (544, 960))
+    _WARMUP_TASKS = ("t2av", "fl2av", "i2av", "l2av", "ref2av")
+
     def __init__(self, config):
         if config.get("task") not in {"t2av", "i2av", "l2av", "fl2av", "ref2av"}:
             raise ValueError("MiniMax-H3 supports t2av/i2av/l2av/fl2av/ref2av")
@@ -73,8 +76,71 @@ def init_modules(self):
         super().init_modules()
         self.run_input_encoder = self._run_input_encoder_local_h3
 
+    @ProfilingContext4DebugL1("Warmup")
     def run_warmup(self):
-        raise NotImplementedError("MiniMax-H3 warmup is not implemented")
+        task = self.config["task"]
+        if task not in self._WARMUP_TASKS:
+            raise NotImplementedError(f"MiniMax-H3 warmup does not support task: {task}")
+
+        for height, width in self._WARMUP_RESOLUTIONS:
+            logger.info(f"Warmup: {height}x{width}")
+            transformer_offloaded = not self.config.get("cpu_offload", False)
+            try:
+                self.scheduler.generator = None
+                self._prepare_warmup_inputs(height, width)
+                self.inputs = self._run_input_encoder_local_h3()
+                self.init_run()
+
+                self.scheduler.step_pre(0)
+                self.model.infer(self.inputs)
+                self.scheduler.step_post()
+                video_rows = self.scheduler.video_latents
+                audio_rows = self.scheduler.audio_latents
+
+                if self.config.get("cpu_offload", False):
+                    self._offload_transformer()
+                    transformer_offloaded = True
+
+                self.run_vae_decoder(video_rows, audio_rows)
+                torch_device_module.synchronize()
+            finally:
+                if self.config.get("cpu_offload", False) and not transformer_offloaded:
+                    with suppress(Exception):
+                        self._offload_transformer()
+                self.clear_warmup_state()
+
+        logger.info("[Warmup] Warmup completed")
+        self._maybe_freeze_gc()
+
+    def _prepare_warmup_inputs(self, height, width):
+        task = self.config["task"]
+        common = {
+            "seed": 0,
+            "prompt": "warmup" if (height, width) == self._WARMUP_RESOLUTIONS[0] else "A cinematic fox walking through a snowy forest.",
+            "target_shape": [height, width],
+            "target_video_length": int(self.config.get("target_video_length", 124)),
+            "return_result_tensor": True,
+        }
+        image = Image.new("RGB", (width, height), color=0)
+        if task == "t2av":
+            self.input_info = T2AVInputInfo(**common)
+        elif task == "i2av":
+            self.input_info = I2AVInputInfo(**common, image_path=image)
+        elif task == "l2av":
+            self.input_info = L2AVInputInfo(**common, last_frame_path=image)
+        elif task == "fl2av":
+            self.input_info = FL2AVInputInfo(**common, image_path=image, last_frame_path=image.copy())
+        else:
+            self.input_info = Ref2AVInputInfo(**common, image_path=image)
+
+    def clear_warmup_state(self):
+        self.scheduler.clear()
+        self.condition_video_latents = []
+        self.condition_audio_latents = []
+        self.keyframe_anchors = ()
+        self.prepared_references = None
+        self.input_info = None
+        self.__dict__.pop("inputs", None)
 
     def init_scheduler(self):
         self.scheduler = MiniMaxH3Scheduler(self.config)

scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+# set path firstly
+lightx2v_path=
+model_path=
+
+export CUDA_VISIBLE_DEVICES=0
+
+# set environment variables
+source ${lightx2v_path}/scripts/base/base.sh
+export DTYPE=BF16
+export SENSITIVE_LAYER_DTYPE=BF16
+
+prompt='A cinematic fox walking through a snowy forest.'
+
+python -m lightx2v.infer \
+--model_cls minimax_h3 \
+--task t2av \
+--model_path $model_path \
+--config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_compile.json \
+--warmup \
+--prompt "$prompt" \
+--save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av10.mp4 \
+--seed 0