diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py
index 7145136..7b4ba8f 100644
--- a/extensions_built_in/diffusion_models/__init__.py
+++ b/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,
]
diff --git a/extensions_built_in/diffusion_models/minimax_h3/__init__.py b/extensions_built_in/diffusion_models/minimax_h3/__init__.py
index 9beb6a8..e1a994c 100644
--- a/extensions_built_in/diffusion_models/minimax_h3/__init__.py
+++ b/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
diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
index 6b0084b..add43d1 100644
--- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
+++ b/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"
diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/packing.py b/extensions_built_in/diffusion_models/minimax_h3/src/packing.py
index c24578d..cb1d272 100644
--- a/extensions_built_in/diffusion_models/minimax_h3/src/packing.py
+++ b/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
):
diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py
index eadcf8d..8d425b7 100644
--- a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py
+++ b/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()
diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py b/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py
index 1352411..306b9e0 100644
--- a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py
+++ b/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
+ )
diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx
index a56d026..53948b5 100644
--- a/extensions_built_in/diffusion_models/ui.tsx
+++ b/extensions_built_in/diffusion_models/ui.tsx
@@ -1472,6 +1472,218 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
),
},
+ {
+ 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: (
+
+ 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.
+
+ How still-image references (dataset control images and sample
+ ctrl images) are shown to the model. Video references always use
+ the video path.
+
+
+ Picture: the native ref2va recipe — a
+ single-frame reference block, shown to Qwen3-VL as a{" "}
+ <Picture i> block, scaled down only.
+
+
+ Static video clip: 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,{" "}
+ <Video k> 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 model_kwargs.image_ref_video_frames{" "}
+ (17n+5). Changing this re-caches text embeddings.
+
+
+ ),
+ },
+ },
+ ],
+ modelNotes: (
+
+
+ FastVideo'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.
+
+
+ Weights load from the{" "}
+
+ Models Folder Path
+ {" "}
+ set in settings. Anything missing is downloaded there from{" "}
+ FastVideo/FastH3-Omni-Comfy 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:
+
+ 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.
+