Merge remote-tracking branch 'origin/aryan/open-h3-regenerate-2k' into aryan/regen2k-nvfp4-decoder

# Conflicts:
#	fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py

docs-site/content/cookbook/minimax-h3/open-h3-regenerate-2k.mdx

@@ -0,0 +1,7 @@
+---
+title: Open-H3-Regenerate-2K
+description: Configure Open-H3-Regenerate-2K serving, which regenerates a 768P MiniMax-H3 draft at 2560x1440.
+mode: frame
+---
+
+<Cookbook modelId="open-h3-regenerate-2k" />

docs-site/lib/cookbook-pages.ts

@@ -133,6 +133,7 @@ const backendModels: Record<string, string> = {
   "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2": "fasth3-preview",
   "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree": "fasth3-4step",
   "FastVideo/FastH3-OmniRef-NVFP4-QAD": "fasth3-omniref",
+  "Efficient-Large-Model/Open-H3-Regenerate-2K": "open-h3-regenerate-2k",
   "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": "fastwan21",
   "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": "wan21-i2v",
   "FastVideo/LTX2-Distilled-Diffusers": "ltx2-distilled",

docs-site/navigation.json

@@ -174,6 +174,10 @@
             "label": "FastH3 4-Step V1",
             "href": "/cookbook/minimax-h3/fasth3-4step"
           },
+          {
+            "label": "Open-H3-Regenerate-2K",
+            "href": "/cookbook/minimax-h3/open-h3-regenerate-2k"
+          },
           {
             "label": "Recipes and setup",
             "href": "/cookbook/minimax-h3"

docs/assets/cookbook-recipes.json

@@ -764,7 +764,7 @@
       "id": "open-h3-regenerate-2k",
       "family": "minimax_h3",
       "stage": "inference",
-      "task": "Video to 2K video",
+      "task": "Video super-resolution (2K regeneration)",
       "label": "Open-H3-Regenerate-2K",
       "summary": "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with NVlabs' three-step refiner. The output is silent.",
       "model": "Efficient-Large-Model/Open-H3-Regenerate-2K",

fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py

@@ -14,6 +14,7 @@ from typing import Any
 
 import torch
 
+from fastvideo import envs
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler
@@ -36,6 +37,13 @@ _BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae"
                     "audio_scheduler")
 
 
+def refuse_adaln_cache() -> None:
+    """Raise when an H3 AdaLN cache is enabled: it bypasses the projection modules that carry the Sharp branches."""
+    if envs.FASTVIDEO_H3_ADALN_CACHE.get() or envs.FASTVIDEO_H3_ADALN_TABLE.get():
+        raise ValueError("Open-H3-Regenerate-2K needs the AdaLN projections: unset FASTVIDEO_H3_ADALN_CACHE and "
+                         "FASTVIDEO_H3_ADALN_TABLE, whose cached modulation skips the 50 Sharp curve branches.")
+
+
 class OpenH3Regenerate2KPipeline(MiniMaxH3Pipeline):
     """MiniMax-H3 T2VA layout with a draft-latent start, fixed sigmas and a silent output."""
 
@@ -49,6 +57,7 @@ class OpenH3Regenerate2KPipeline(MiniMaxH3Pipeline):
         return [f"{BASE_MODEL_DIR}/model_index.json", *components, "lora/**", "latent_upsampler/*.safetensors"]
 
     def _load_config(self, model_path: str) -> dict[str, Any]:
+        refuse_adaln_cache()
         revision = getattr(self.fastvideo_args, "revision", None)
         self.model_path = maybe_download_model(model_path,
                                                revision=revision,

fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py

@@ -60,8 +60,21 @@ class Regenerate2KDenoisingStage(PipelineStage):
         device = get_local_torch_device()
         sigmas = refine_sigmas(batch.num_inference_steps, self.start_sigma)
         state = batch.latents.to(device)
-        with self._attention_kernel(device):
-            state = self._euler_loop(batch, layout, state, sigmas, device, transformer)
+        layerwise = bool(getattr(fastvideo_args, "dit_layerwise_offload", False))
+        full_cpu_offload = (bool(getattr(fastvideo_args, "dit_cpu_offload", False)) and not layerwise
+                            and not bool(getattr(fastvideo_args, "use_fsdp_inference", False)))
+        try:
+            if full_cpu_offload:
+                transformer.to(device)
+            with self._attention_kernel(device):
+                state = self._euler_loop(batch, layout, state, sigmas, device, transformer)
+        finally:
+            if layerwise:
+                manager = getattr(transformer, "_layerwise_offload_manager", None)
+                if manager is not None and getattr(manager, "enabled", False):
+                    manager.release_all()
+            if full_cpu_offload:
+                transformer.to("cpu")
         batch.latents = state
         return batch
 

fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py

@@ -243,7 +243,7 @@ def test_stage_attaches_the_branches_to_each_materialized_lazy_transformer(monke
     assert second.targets[-1] is not loaded_first
 
 
-def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None):
+def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None, fastvideo_args=None):
     from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence
     from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY
     from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising
@@ -260,7 +260,7 @@ def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None):
     transformer = transformer if transformer is not None else _ConstantVelocityTransformer(velocity)
     curve = curve if curve is not None else _RecordingCurve()
     stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None)
-    return stage.forward(batch, None), transformer, curve
+    return stage.forward(batch, fastvideo_args), transformer, curve
 
 
 def test_euler_steps_integrate_the_constant_velocity_over_sigma_zero_to_start(monkeypatch):
@@ -346,3 +346,45 @@ def test_generate_dotted_override_sets_the_nvfp4_decoder_option(tmp_path):
     assert config.generator.pipeline.experimental["h3_vae_nvfp4_decoder"] == "/tmp/decoder-nvfp4.safetensors"
     assert generator_config_to_fastvideo_args(config.generator).h3_vae_nvfp4_decoder == "/tmp/decoder-nvfp4.safetensors"
 
+
+class _OffloadedTransformer(_ConstantVelocityTransformer):
+
+    def __init__(self, velocity, fail=False):
+        super().__init__(velocity)
+        self.fail = fail
+        self.moves = []
+
+    def to(self, device):
+        self.moves.append(str(device))
+        return self
+
+    def __call__(self, **kwargs):
+        if self.fail:
+            raise RuntimeError("forward failed")
+        return super().__call__(**kwargs)
+
+
+@pytest.mark.parametrize("fail", [False, True])
+def test_full_cpu_offload_moves_the_transformer_back_to_cpu_after_the_refine(monkeypatch, fail):
+    from types import SimpleNamespace
+
+    transformer = _OffloadedTransformer(0.5, fail=fail)
+    args = SimpleNamespace(dit_cpu_offload=True, dit_layerwise_offload=False, use_fsdp_inference=False)
+    if fail:
+        with pytest.raises(RuntimeError, match="forward failed"):
+            _denoise(monkeypatch, steps=2, transformer=transformer, fastvideo_args=args)
+    else:
+        _denoise(monkeypatch, steps=2, transformer=transformer, fastvideo_args=args)
+
+    assert transformer.moves == ["cpu", "cpu"]  # to the stage device (CPU in this test), then back to the host
+
+
+@pytest.mark.parametrize("variable, value", [("FASTVIDEO_H3_ADALN_CACHE", True),
+                                              ("FASTVIDEO_H3_ADALN_TABLE", "/tmp/adaln_tables.pt")])
+def test_pipeline_refuses_an_adaln_cache_that_would_skip_the_sharp_branches(variable, value):
+    from fastvideo import envs
+    from fastvideo.pipelines.basic.open_h3_regenerate_2k.open_h3_regenerate_2k_pipeline import refuse_adaln_cache
+
+    refuse_adaln_cache()
+    with getattr(envs, variable).override(value), pytest.raises(ValueError, match="Sharp curve branches"):
+        refuse_adaln_cache()