Add support to train Fast H3 Omni with training adapter

extensions_built_in/diffusion_models/__init__.py

@@ -22,7 +22,7 @@ from .krea2 import Krea2Model
 from .iris3b import Iris3BModel
 from .boogu_image import BooguImageModel, BooguImageEditModel
 from .mageflow import MageFlowModel, MageFlowEditModel
-from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel
+from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel, MinimaxH3OmniModel
 from .flux3 import Flux3Model
 from .moss_image import MossImageModel
 
@@ -67,6 +67,7 @@ AI_TOOLKIT_MODELS = [
     MinimaxH3Model,
     MinimaxH3Ref2VAModel,
     MinimaxH3FastModel,
+    MinimaxH3OmniModel,
     Flux3Model,
     MossImageModel,
 ]

extensions_built_in/diffusion_models/minimax_h3/__init__.py

@@ -1 +1 @@
-from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel
+from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel, MinimaxH3OmniModel

extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py

@@ -120,10 +120,13 @@ COMFY_FILES = {
     "video_vae": "vae/minimax_h3_video_vae_fp16.safetensors",
     "audio_vae": "vae/minimax_h3_audio_vae_fp32.safetensors",
     "dit_fasth3": "diffusion_models/fastvideo_fasth3_8step_v2_pruned_int8_convrot.safetensors",
+    "dit_fasth3_omni": "diffusion_models/fasth3_omniref_8step_pruned_int8_convrot.safetensors",
 }
 # FastH3 8-Step V2 (FastVideo DMD2 + VSA distill) comfy repack: DiT plus the
 # shared text encoder and VAEs at the same repo-relative paths
 FASTH3_REPO = "FastVideo/FastVideo-FastH3-Comfy"
+# FastH3 Omni (8-step ref2va distill, 32 PDD heads) comfy repack, same layout
+FASTH3_OMNI_REPO = "FastVideo/FastH3-Omni-Comfy"
 # tokenizer/processor/text-encoder config come from the original repo (tiny files)
 ORIGINAL_REPO = "MiniMaxAI/MiniMax-H3"
 
@@ -955,6 +958,7 @@ class MinimaxH3Model(BaseModel):
             text_indices=text_indices.to(device),
             # target-video token grid (patch 1x2x2); consumed only by VSA models
             vsa_video_grid=(t_lat, h_lat // 2, w_lat // 2),
+            pdd_weights=self._pdd_weights(sigma_v),
         )
 
         if num_cond_audio > 0:
@@ -975,6 +979,19 @@ class MinimaxH3Model(BaseModel):
             )
         return -noise_pred
 
+    def _pdd_weights(self, sigma_v, sigma_v_next=None):
+        """Head weights for PDD checkpoints (None otherwise). Training passes a
+        single sigma: the one 1/n interval containing it."""
+        n = unwrap_model(self.model).params.pdd_steps
+        if n <= 1:
+            return None
+        return packing.build_pdd_weights(
+            sigma_v.flatten().tolist(),
+            n,
+            self.video_sigma_shift,
+            sigma_v_next=None if sigma_v_next is None else sigma_v_next.flatten().tolist(),
+        )
+
     def get_loss_target(self, *args, **kwargs):
         noise = kwargs.get("noise")
         batch = kwargs.get("batch")
@@ -1618,3 +1635,59 @@ class MinimaxH3FastModel(MinimaxH3Model):
     def text_embedding_space_version(self):
         # same Qwen3-VL presentation as the base model: share the embed cache
         return "minimax_h3"
+
+
+class MinimaxH3OmniModel(MinimaxH3Ref2VAModel):
+    """FastH3 Omni: FastVideo's 8-step ref2va distill. References work exactly
+    like ``minimax_h3_ref2va`` (control images/videos as reference blocks).
+
+    The DiT carries 32 PDD output heads per stream (see
+    MiniMaxH3TransformerParams.pdd_steps): sampling blends the heads each
+    step spans, training uses the head of the interval containing its sigma.
+    Video/audio shifts stay 12/3.
+
+    It was distilled with FastVideo's VSA (tile 128, ``p2_multi_region``
+    reference policy), which src/vsa.py does not implement, so attention runs
+    dense like native ComfyUI. ``model_kwargs.vsa`` (default false) opts into
+    the (4,4,4) tile VSA with references as dense prefix, at
+    ``vsa_sparsity`` (default 0.9).
+
+    Sample with 8 steps and guidance_scale 1.
+    """
+
+    arch = "minimax_h3_omni"
+    comfy_repo = FASTH3_OMNI_REPO
+    # trained ladder: [999, 874, ..., 125] -> PDD spans [0, 4, ..., 32]
+    t1000_sample_ladder = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        kw = self.model_config.model_kwargs
+        self.vsa_sparsity: Optional[float] = None
+        if bool(kw.get("vsa", False)):
+            self.vsa_sparsity = float(kw.get("vsa_sparsity", 0.9))
+
+    def _dit_component(self) -> str:
+        return "dit_fasth3_omni"
+
+    def _load_transformer(self) -> MiniMaxH3Transformer:
+        transformer = super()._load_transformer()
+        if transformer.params.pdd_steps <= 1:
+            raise ValueError(
+                "minimax_h3_omni needs the FastH3 Omni checkpoint with stacked "
+                "PDD output heads; this file has a single head. Use arch "
+                "minimax_h3_ref2va for the base ref2va checkpoints."
+            )
+        transformer.vsa_sparsity = self.vsa_sparsity
+        return transformer
+
+    def get_base_model_version(self):
+        return "minimax_h3_omni"
+
+    @property
+    def text_embedding_space_version(self):
+        # same Qwen3-VL reference presentation as ref2va: share its embed cache
+        n = self._image_ref_video_frames()
+        if n:
+            return f"minimax_h3_ref2va:img_as_vid{n}"
+        return "minimax_h3_ref2va"

extensions_built_in/diffusion_models/minimax_h3/src/packing.py

@@ -536,6 +536,49 @@ def shift_sigma(sigma, shift: float):
     return shift * sigma / (1.0 + (shift - 1.0) * sigma)
 
 
+def pdd_span(sigma_v: float, n: int, video_shift: float, sigma_v_next=None):
+    """Head interval [start, stop) on the unshifted n-interval grid. A sampling
+    step spans sigma -> sigma_next (ComfyUI's rounding); a single training
+    sigma uses the one interval containing it."""
+    base_t = 1.0 - float(shift_sigma(sigma_v, 1.0 / video_shift))
+    if sigma_v_next is None:
+        start = min(max(int(math.floor(base_t * n)), 0), n - 1)
+        return start, start + 1
+    base_t_next = 1.0 - float(shift_sigma(sigma_v_next, 1.0 / video_shift))
+    start = min(round(base_t * n), n - 1)
+    stop = max(round(base_t_next * n), start + 1)
+    return start, stop
+
+
+def pdd_head_weights(start: int, stop: int, n: int, flow_shift: float):
+    """(n,) weights over the stacked heads: head 0 (the full head) at 1, the
+    offset heads in [max(start, 1), stop) at their normalized dt share."""
+    grid = torch.linspace(1.0, 0.0, n + 1, dtype=torch.float64)
+    dt = (1.0 - shift_sigma(grid, flow_shift)).diff()[start:stop]
+    w = torch.zeros(n, dtype=torch.float64)
+    w[0] = 1.0
+    first = max(start, 1)
+    w[first:stop] = (dt / dt.sum())[first - start :]
+    return w
+
+
+def build_pdd_weights(
+    sigma_v,  # (B,) video sigmas
+    n: int,
+    video_shift: float,
+    audio_shift: float = AUDIO_SIGMA_SHIFT,
+    sigma_v_next=None,  # (B,) or None for single-sigma training
+):
+    """Per-sample (video, audio) head weights, each (B, n) float32."""
+    wv, wa = [], []
+    for i in range(len(sigma_v)):
+        nxt = None if sigma_v_next is None else float(sigma_v_next[i])
+        start, stop = pdd_span(float(sigma_v[i]), n, video_shift, nxt)
+        wv.append(pdd_head_weights(start, stop, n, video_shift))
+        wa.append(pdd_head_weights(start, stop, n, audio_shift))
+    return torch.stack(wv).float(), torch.stack(wa).float()
+
+
 def remap_sigma(
     sigma, from_shift: float = VIDEO_SIGMA_SHIFT, to_shift: float = AUDIO_SIGMA_SHIFT
 ):

extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py

@@ -245,6 +245,7 @@ class MiniMaxH3Pipeline:
                 text_indices=text_indices,
                 # target-video token grid (patch 1x2x2); consumed only by VSA models
                 vsa_video_grid=(t_lat, h_lat // 2, w_lat // 2),
+                pdd_weights=model._pdd_weights(sv, sv_next),
             )
             v_video = video_pred[:, num_cond:].float()
             v_audio = audio_pred[:, layout.num_condition_audio_rows :].float()

extensions_built_in/diffusion_models/minimax_h3/src/transformer.py

@@ -71,6 +71,10 @@ class MiniMaxH3TransformerParams:
     # and add a bias to the block AdaLN linears, which the original weights
     # lack (the final layer's AdaLN carries a bias in both variants).
     adaln_t_table_size: Optional[int] = None
+    # PDD checkpoints (FastH3 Omni) stack pdd_steps output heads per stream on
+    # the unshifted [1, 0] grid: block 0 is a full head, later blocks are
+    # offsets from it; a step uses the dt-weighted blend of the heads it spans
+    pdd_steps: int = 1
 
     @property
     def adaln_apply_silu(self) -> bool:
@@ -368,14 +372,22 @@ class MiniMaxH3FinalLayer(nn.Module):
             apply_silu=p.adaln_apply_silu,
             bias=True,
         )
-        self.video_out = nn.Linear(p.hidden_size, video_patch_dim, bias=True)
-        self.audio_out = nn.Linear(p.hidden_size, p.audio_latents_dim, bias=True)
+        self.pdd_steps = p.pdd_steps
+        self.video_out = nn.Linear(
+            p.hidden_size, video_patch_dim * p.pdd_steps, bias=True
+        )
+        self.audio_out = nn.Linear(
+            p.hidden_size, p.audio_latents_dim * p.pdd_steps, bias=True
+        )
 
     def forward(
         self,
         x: torch.Tensor,  # (B, S, hidden)
         temb: torch.Tensor,  # (M, time_embed_dim)
         timestep_indices: torch.Tensor,  # (B, S) long into the M rows
+        video_indices: torch.Tensor,  # (Nv,) long
+        audio_indices: torch.Tensor,  # (Na,) long
+        pdd_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ):
         shift, scale = self.adaln_proj(temb)
         dt = x.dtype
@@ -383,7 +395,23 @@ class MiniMaxH3FinalLayer(nn.Module):
             timestep_indices
         ].to(dt)
         h = h.to(self.video_out.weight.dtype)
-        return self.video_out(h), self.audio_out(h)
+        video = self.video_out(h.index_select(1, video_indices))
+        audio = self.audio_out(h.index_select(1, audio_indices))
+        if self.pdd_steps > 1:
+            if pdd_weights is None:
+                raise ValueError("PDD heads need pdd_weights for the current step")
+            # heads run through the Linear's own forward so LoRA hooks apply
+            video = _blend_pdd_heads(video, pdd_weights[0], self.pdd_steps)
+            audio = _blend_pdd_heads(audio, pdd_weights[1], self.pdd_steps)
+        return video, audio
+
+
+def _blend_pdd_heads(out: torch.Tensor, weights: torch.Tensor, n: int):
+    """(B, N, n * d) stacked head outputs, (B, n) weights -> (B, N, d)."""
+    b, rows, _ = out.shape
+    return torch.einsum(
+        "brkd,bk->brd", out.view(b, rows, n, -1), weights.to(out.device, out.dtype)
+    )
 
 
 class MiniMaxH3Transformer(nn.Module, OstrisModelMixin):
@@ -399,6 +427,12 @@ class MiniMaxH3Transformer(nn.Module, OstrisModelMixin):
             params.adaln_t_table_size = table.shape[0]
             params.time_embed_dim = table.shape[1]
         params.gate_compress = "blocks.0.attn.to_gate_compress.weight" in state_dict
+        video_out = state_dict.get("final_layer.video_out.weight", None)
+        if video_out is not None:
+            p = params.patch_size
+            params.pdd_steps = video_out.shape[0] // (
+                params.latents_dim * p[0] * p[1] * p[2]
+            )
         return params
 
     @classmethod
@@ -504,6 +538,9 @@ class MiniMaxH3Transformer(nn.Module, OstrisModelMixin):
         vsa_video_grid: Optional[
             Tuple[int, int, int]
         ] = None,  # target-video token grid (t, h, w)
+        # PDD checkpoints only: per-sample (video, audio) head weights, each
+        # (B, pdd_steps), from packing.build_pdd_weights
+        pdd_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ):
         """Returns (video_out (B, Nv, 96), audio_out (B, Na, 32)) — the
         data-ward velocity ``clean - noise`` for every row, in input order.
@@ -582,7 +619,6 @@ class MiniMaxH3Transformer(nn.Module, OstrisModelMixin):
             else:
                 x = block(x, temb, adaln_indices, rotary_emb, attn_mask, vsa_ctx)
 
-        video_all, audio_all = self.final_layer(x, temb, inverse)
-        video_out = video_all.index_select(1, video_indices)
-        audio_out = audio_all.index_select(1, audio_indices)
-        return video_out, audio_out
+        return self.final_layer(
+            x, temb, inverse, video_indices, audio_indices, pdd_weights
+        )

extensions_built_in/diffusion_models/ui.tsx

@@ -1472,6 +1472,218 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       </div>
     ),
   },
+  {
+    name: "minimax_h3_omni",
+    label: "FastH3 Omni",
+    group: "video",
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "FastVideo/FastH3-Omni-Comfy",
+        defaultNameOrPath,
+      ],
+      // pre-quantized weights: matching qtypes keep the load unchanged
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.qtype_te": ["nvfp4", "qfloat8"],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].sample.sampler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.cache_text_embeddings": [true, false],
+      "config.process[0].model.assistant_lora_path": [
+        "ostris/minimax_h3_training_adapter/fastvideo_fasth3_omni_training_adapter_v1.safetensors",
+        undefined,
+      ],
+      "config.process[0].network.linear": [16, defaultLinearRank],
+      "config.process[0].network.linear_alpha": [16, defaultLinearRank],
+      "config.process[0].network.network_kwargs.ignore_if_contains": [
+        ["adaln_proj"],
+        [],
+      ],
+      "config.process[0].sample.num_frames": [107, 1],
+      "config.process[0].sample.fps": [24, 1],
+      "config.process[0].sample.width": [768, 1024],
+      "config.process[0].sample.height": [768, 1024],
+      "config.process[0].sample.guidance_scale": [1, 4],
+      "config.process[0].sample.sample_steps": [8, 25],
+      "config.process[0].train.audio_loss_multiplier": [1.0, undefined],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      "config.process[0].datasets[x].do_audio": [true, undefined],
+      "config.process[0].datasets[x].cache_latents_to_disk": [true, false],
+      "config.process[0].datasets[x].fps": [24, undefined],
+      "config.process[0].datasets[x].num_frames": [39, undefined],
+      "config.process[0].datasets[x].auto_frame_count": [true, undefined],
+    },
+    disableSections: ["network.conv"],
+    additionalSections: [
+      "sample.multi_ctrl_imgs",
+      "datasets.multi_control_paths",
+      "datasets.num_frames",
+      "model.layer_offloading",
+      "model.low_vram",
+      "datasets.do_audio",
+      "datasets.audio_normalize",
+      "datasets.audio_preserve_pitch",
+      "train.audio_loss_multiplier",
+      "datasets.auto_frame_count",
+      "model.assistant_lora_path",
+    ],
+    customModelSelectOptions: [
+      {
+        label: "Distillation Handling Method",
+        options: [
+          { value: "ta", label: "Training Adapter (default)" },
+          { value: "dopsd", label: "D-OPSD" },
+          { value: "none", label: "None" },
+        ],
+        getValue: (config: JobConfig) => {
+          if (config?.config?.process?.[0]?.model?.model_kwargs?.dopsd) {
+            return "dopsd";
+          }
+          const assistantLoraPath =
+            config?.config?.process?.[0]?.model?.assistant_lora_path;
+          return assistantLoraPath && assistantLoraPath.trim() !== ""
+            ? "ta"
+            : "none";
+        },
+        onChange: (
+          value: string,
+          config: JobConfig,
+          setJobConfig: (value: any, key: string) => void,
+        ) => {
+          const kwargs = {
+            ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}),
+          };
+          if (value === "dopsd") {
+            kwargs.dopsd = true;
+            kwargs.dopsd_bleed_strength = 1.0;
+          } else {
+            delete kwargs.dopsd;
+            delete kwargs.dopsd_bleed_strength;
+          }
+          setJobConfig(kwargs, "config.process[0].model.model_kwargs");
+          setJobConfig(
+            value === "ta"
+              ? "ostris/minimax_h3_training_adapter/fastvideo_fasth3_omni_training_adapter_v1.safetensors"
+              : undefined,
+            "config.process[0].model.assistant_lora_path",
+          );
+        },
+        doc: {
+          title: "FastH3 Omni Distillation Handling",
+          description: (
+            <div>
+              FastH3 Omni is a step-distilled model, so training on it directly
+              will make the distillation break down. The Training Adapter (the
+              MiniMax-H3 Ref2V adapter) absorbs that drift during training and
+              is removed at inference, but can still break down over a long
+              run. D-OPSD instead self-distills: a no-grad teacher pass sees the
+              training target as its own reference and its prediction becomes
+              the training target for a reference-free pass, baking the
+              reference into your trigger word (or the caption itself when no
+              trigger word is set). Re-caches latents with pixel tensors.
+            </div>
+          ),
+        },
+      },
+      {
+        label: "Image Reference Presentation",
+        options: [
+          { value: "picture", label: "Picture (default)" },
+          { value: "video", label: "Static video clip" },
+        ],
+        getValue: (config: JobConfig) => {
+          return config?.config?.process?.[0]?.model?.model_kwargs
+            ?.image_refs_as_video
+            ? "video"
+            : "picture";
+        },
+        onChange: (
+          value: string,
+          config: JobConfig,
+          setJobConfig: (value: any, key: string) => void,
+        ) => {
+          const kwargs = {
+            ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}),
+          };
+          if (value === "video") {
+            kwargs.image_refs_as_video = true;
+          } else {
+            delete kwargs.image_refs_as_video;
+            delete kwargs.image_ref_video_frames;
+          }
+          setJobConfig(kwargs, "config.process[0].model.model_kwargs");
+        },
+        doc: {
+          title: "MiniMax-H3 Image Reference Presentation",
+          description: (
+            <div className="space-y-2">
+              <p>
+                How still-image references (dataset control images and sample
+                ctrl images) are shown to the model. Video references always use
+                the video path.
+              </p>
+              <p>
+                <strong>Picture</strong>: the native ref2va recipe — a
+                single-frame reference block, shown to Qwen3-VL as a{" "}
+                <code>&lt;Picture i&gt;</code> block, scaled down only.
+              </p>
+              <p>
+                <strong>Static video clip</strong>: the image is held for 5
+                frames (2 latent frames) and routed through the exact path a
+                reference VIDEO takes — video sizing (matched to the target's
+                pixel area), multi-frame reference block,{" "}
+                <code>&lt;Video k&gt;</code> timestamped presentation. Use this
+                when training on image references but sampling with video
+                references, so the LoRA learns the pathway it will be used
+                through. Adds a handful of rows per reference. Frame count is
+                adjustable with <code>model_kwargs.image_ref_video_frames</code>{" "}
+                (17n+5). Changing this re-caches text embeddings.
+              </p>
+            </div>
+          ),
+        },
+      },
+    ],
+    modelNotes: (
+      <div className="space-y-2">
+        <p>
+          FastVideo&apos;s FastH3 Omni: MiniMax-H3 Ref2V distilled to 8 steps.
+          References work exactly like MiniMax-H3 Ref2V: control images and
+          videos condition the output as subject/style references (never as a
+          first frame), from the dataset control path(s) in training and the
+          sample ctrl images in sampling.
+        </p>
+        <p>
+          Weights load from the{" "}
+          <Link href="/settings" className="text-blue-400 hover:underline">
+            Models Folder Path
+          </Link>{" "}
+          set in settings. Anything missing is downloaded there from{" "}
+          <code>FastVideo/FastH3-Omni-Comfy</code> on first load. The text
+          encoder and VAEs are the same files MiniMax-H3 uses, so an existing
+          MiniMax-H3 setup only needs the new DiT. Files used:
+        </p>
+        <pre className="bg-gray-900 border border-gray-700 rounded-lg p-3 text-xs overflow-x-auto">
+          <code>{`<MODELS_PATH>/
+├── diffusion_models/
+│   └── fasth3_omniref_8step_pruned_int8_convrot.safetensors
+├── text_encoders/
+│   └── qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors
+└── vae/
+    ├── minimax_h3_video_vae_fp16.safetensors
+    └── minimax_h3_audio_vae_fp32.safetensors`}</code>
+        </pre>
+        <p>
+          Sample with 8 steps and guidance scale 1 (the trained schedule).
+          Attention runs dense, like native ComfyUI. Video is fixed 24 fps and
+          frame counts snap down to the 17n+5 grid.
+        </p>
+      </div>
+    ),
+  },
   {
     name: "minimax_h3_vsa",
     label: "FastH3 8-Step V2",

extensions_built_in/inference_engine/latent_preview.py

@@ -47,6 +47,7 @@ ARCH_LATENT_FORMAT = {
     "minimax_h3": "MiniMaxH3Video",
     "minimax_h3_ref2va": "MiniMaxH3Video",
     "minimax_h3_vsa": "MiniMaxH3Video",
+    "minimax_h3_omni": "MiniMaxH3Video",
     "anima": "Wan21",
     "cogview4": "Flux",
     "prx_pixel": None,  # pixel space