From a776394353c0c6e0a65c838e659f776b972e95a2 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 24 Sep 2026 18:40:17 -0600 Subject: [PATCH] Add support to train on Fast H3 v2 8 step directly with training adapter --- .../diffusion_models/minimax_h3/minimax_h3.py | 103 +++++--------- .../minimax_h3/src/packing.py | 2 +- .../minimax_h3/src/pipeline.py | 5 +- .../diffusion_models/minimax_h3/src/vsa.py | 2 +- extensions_built_in/diffusion_models/ui.tsx | 129 ++++++++++++++++++ 5 files changed, 167 insertions(+), 74 deletions(-) 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 248fa14..6b0084b 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -119,15 +119,11 @@ COMFY_FILES = { "text_encoder": "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "video_vae": "vae/minimax_h3_video_vae_fp16.safetensors", "audio_vae": "vae/minimax_h3_audio_vae_fp32.safetensors", - "dit_fasth3": "diffusion_models/minimax_h3_fasth3_preview_v0.2_int8_convrot.safetensors", + "dit_fasth3": "diffusion_models/fastvideo_fasth3_8step_v2_pruned_int8_convrot.safetensors", } -# FastH3 (FastVideo 4-step VSA distill) int8-convrot repack, produced by -# scripts/convert_minimax_h2_fastvideo.py from the FastVideo diffusers repo. -# Hub fallback repo for the file (flat at the repo root, downloaded into -# MODELS_PATH/diffusion_models/); not published there yet — until it is, the -# file must exist locally or be built with the converter. -FASTH3_REPO = "Kijai/MiniMax-H3-experimental" -FASTH3_SOURCE_REPO = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2" +# 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" # tokenizer/processor/text-encoder config come from the original repo (tiny files) ORIGINAL_REPO = "MiniMaxAI/MiniMax-H3" @@ -180,6 +176,10 @@ class MiniMaxH3VaeBundle(torch.nn.Module): class MinimaxH3Model(BaseModel): arch = "minimax_h3" + # video flow shift; audio stays on AUDIO_SIGMA_SHIFT, remapped from this + video_sigma_shift = packing.VIDEO_SIGMA_SHIFT + # hub repo for missing files when name_or_path isn't an org/repo id + comfy_repo = COMFY_REPO use_old_lokr_format = False def __init__( @@ -216,9 +216,11 @@ class MinimaxH3Model(BaseModel): self.model_config.model_kwargs.get("max_text_length", 512) ) - @staticmethod - def get_train_scheduler(): - return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) + @classmethod + def get_train_scheduler(cls): + return CustomFlowMatchEulerDiscreteScheduler( + **{**scheduler_config, "shift": cls.video_sigma_shift} + ) def get_bucket_divisibility(self): # 16x VAE spatial compression * 2x2 transformer patch @@ -256,7 +258,7 @@ class MinimaxH3Model(BaseModel): ) return resolve_comfy_file( COMFY_FILES[component], - repo_id=repo_id_from_name_or_path(name_or_path, COMFY_REPO), + repo_id=repo_id_from_name_or_path(name_or_path, self.comfy_repo), override_path=self.model_config.model_kwargs.get(f"{component}_path", None), extra_roots=extra_roots, status_fn=self.print_and_status_update, @@ -515,7 +517,7 @@ class MinimaxH3Model(BaseModel): vae_bundle = self._load_vaes() vae_bundle.to(self.vae_device_torch) - self.noise_scheduler = MinimaxH3Model.get_train_scheduler() + self.noise_scheduler = self.get_train_scheduler() self.vae = vae_bundle self.text_encoder = text_encoder self.tokenizer = tokenizer @@ -788,7 +790,9 @@ class MinimaxH3Model(BaseModel): sigma_v = sigma_v.unsqueeze(0) if sigma_v.shape[0] != batch_size: sigma_v = sigma_v.expand(batch_size) - sigma_a = remap_sigma(sigma_v) + sigma_a = remap_sigma( + sigma_v, self.video_sigma_shift, packing.AUDIO_SIGMA_SHIFT + ) t_v = 1.0 - sigma_v t_a = 1.0 - sigma_a @@ -1552,29 +1556,31 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): class MinimaxH3FastModel(MinimaxH3Model): - """FastH3: FastVideo's DMD2-distilled 4-step MiniMax-H3 preview, trained - with VSA (Video Sparse Attention) at 90% sparsity. Text-to-video(+audio) - only — no first-frame or reference conditioning. + """FastH3 8-Step V2: FastVideo's DMD2-distilled 8-step MiniMax-H3, trained + with VSA (Video Sparse Attention) at 80% sparsity and a video flow shift of + 10 (audio stays at 3). Text-to-video(+audio) only — no first-frame or + reference conditioning. - The DiT is the int8-ConvRot repack of - ``FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2`` (built with - scripts/convert_minimax_h2_fastvideo.py): the pruned comfy layout plus a - ``blocks.N.attn.to_gate_compress`` linear per block. Attention runs the - trained VSA policy — pooled 64-token (4,4,4) tile top-k block-sparse + Weights come from ``FASTH3_REPO``. The DiT is its int8-ConvRot repack: + the pruned comfy layout plus a ``blocks.N.attn.to_gate_compress`` linear + per block. Attention runs + the trained VSA policy — pooled 64-token (4,4,4) tile top-k block-sparse attention with the gated compression branch — via FlexAttention (src/vsa.py), in training and sampling alike. Text encoder and VAEs are - shared with the base model. + the base model's files, repacked in the same repo. ``model_kwargs``: - ``vsa`` (default true): false runs dense attention with the gate branch off, matching FastVideo's dense LoRA-preview mode - - ``vsa_sparsity`` (default 0.9, the checkpoint's trained policy) + - ``vsa_sparsity`` (default 0.8, the checkpoint's trained policy) - Sample with 4 steps (the distilled schedule) and guidance_scale 1. + Sample with 8 steps (the distilled schedule) and guidance_scale 1. """ arch = "minimax_h3_vsa" - # FastH3 samples on its trained ladder: [999, 749, 500, 250] on the + video_sigma_shift = 10.0 + comfy_repo = FASTH3_REPO + # FastH3 samples on its trained ladder: [999, 874, ..., 250, 125] on the # shared 1000-step grid, each scheduler applying its own shift t1000_sample_ladder = True @@ -1583,54 +1589,11 @@ class MinimaxH3FastModel(MinimaxH3Model): kw = self.model_config.model_kwargs self.vsa_sparsity: Optional[float] = None if bool(kw.get("vsa", True)): - self.vsa_sparsity = float(kw.get("vsa_sparsity", 0.9)) + self.vsa_sparsity = float(kw.get("vsa_sparsity", 0.8)) def _dit_component(self) -> str: return "dit_fasth3" - def _resolve_comfy_file(self, component: str) -> str: - if component != "dit_fasth3": - return super()._resolve_comfy_file(component) - # local search at the comfy-layout locations first; the hub file sits - # at the FastH3 repo's root, so the download is done directly here - name_or_path = self.model_config.name_or_path - extra_roots = ( - [name_or_path] if name_or_path and os.path.isdir(name_or_path) else [] - ) - rel_path = COMFY_FILES[component] - found = resolve_comfy_file( - rel_path, - repo_id=FASTH3_REPO, - override_path=self.model_config.model_kwargs.get(f"{component}_path", None), - extra_roots=extra_roots, - status_fn=self.print_and_status_update, - local_only=True, - ) - if found is not None: - return found - import huggingface_hub - - target_dir = os.path.join(MODELS_PATH, "diffusion_models") - os.makedirs(target_dir, exist_ok=True) - self.print_and_status_update( - f"Downloading {os.path.basename(rel_path)} from {FASTH3_REPO}" - ) - try: - return huggingface_hub.hf_hub_download( - repo_id=FASTH3_REPO, - filename=os.path.basename(rel_path), - local_dir=target_dir, - ) - except Exception as e: - raise FileNotFoundError( - f"{os.path.basename(rel_path)} was not found locally or on " - f"{FASTH3_REPO}. Build it from the FastVideo release with:\n" - f" python scripts/convert_minimax_h2_fastvideo.py " - f"{FASTH3_SOURCE_REPO} {os.path.join(target_dir, os.path.basename(rel_path))}\n" - f"or point model_kwargs.dit_fasth3_path at an existing " - f"FastH3 int8-convrot file." - ) from e - def _load_transformer(self) -> MiniMaxH3Transformer: transformer = super()._load_transformer() if not transformer.params.gate_compress: 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 0938a14..c24578d 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/packing.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/packing.py @@ -556,7 +556,7 @@ def build_sigma_schedule( `steps`; we don't, so sample_steps means model evals). ``t1000_ladder`` uses FastH3's trained ladder instead: rounded indices on - the shared 1000-step grid ([999, 749, 500, 250] at 4 steps), one forward + the shared 1000-step grid ([999, 874, ..., 250, 125] at 8 steps), one forward per entry, each scheduler applying its own shift.""" if t1000_ladder: base = ( 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 3e48c56..eadcf8d 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py @@ -200,14 +200,15 @@ class MiniMaxH3Pipeline: audio_rows = pack_audio_latents(audio_noise) # (1, 2*A, 32) # --- schedules ----------------------------------------------------- + video_shift = getattr(model, "video_sigma_shift", VIDEO_SIGMA_SHIFT) sigmas_v = build_sigma_schedule( num_inference_steps, - VIDEO_SIGMA_SHIFT, + video_shift, t1000_ladder=getattr(model, "t1000_sample_ladder", False), ).to(device) # the audio schedule follows the video grid through the closed-form # shift remap so both streams sit at the same underlying position - sigmas_a = remap_sigma(sigmas_v, VIDEO_SIGMA_SHIFT, AUDIO_SIGMA_SHIFT) + sigmas_a = remap_sigma(sigmas_v, video_shift, AUDIO_SIGMA_SHIFT) position_ids = layout.position_ids[None].to(device) tags = layout.token_tags[None].to(device) diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/vsa.py b/extensions_built_in/diffusion_models/minimax_h3/src/vsa.py index 7fdc46e..5ea609a 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/vsa.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/vsa.py @@ -14,7 +14,7 @@ a 64-token BlockMask; the compression branch adds ``gate * softmax(scores) @ v_pooled`` per tile, gate from ``to_gate_compress``. -The FastH3 4-step checkpoints were distilled WITH this policy (sparsity 0.9, +The FastH3 8-step V2 checkpoint was distilled WITH this policy (sparsity 0.8, tile 64 is the trained geometry), so training and sampling both run it, at every sequence length. """ diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 109e892..8c8c6cc 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1407,6 +1407,135 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ ), }, + { + name: "minimax_h3_vsa", + label: "FastH3 8-Step V2", + group: "video", + isVideoModel: true, + defaults: { + // default updates when [selected, unselected] in the UI + "config.process[0].model.name_or_path": [ + "FastVideo/FastVideo-FastH3-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_8step_v2_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: [ + "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: "none", label: "None" }, + ], + getValue: (config: JobConfig) => { + 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, + ) => { + setJobConfig( + value === "ta" + ? "ostris/minimax_h3_training_adapter/minimax_h3_training_adapter_v1.safetensors" + : undefined, + "config.process[0].model.assistant_lora_path", + ); + }, + doc: { + title: "FastH3 Distillation Handling", + description: ( +
+ FastH3 is a step-distilled model, so training on it directly will + make the distillation break down. The Training Adapter (the + MiniMax-H3 adapter) absorbs that drift during training and is + removed at inference, but can still break down over a long run. +
+ ), + }, + }, + ], + modelNotes: ( +
+

+ FastVideo's FastH3 8-Step V2: MiniMax-H3 DMD2-distilled to 8 steps + with Video Sparse Attention (80% sparsity). Text-to-video with joint + audio only — no first-frame or reference conditioning. +

+

+ Weights load from the{" "} + + Models Folder Path + {" "} + set in settings. Anything missing is downloaded there from{" "} + FastVideo/FastVideo-FastH3-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: +

+ 
+          {`/
+├── diffusion_models/
+│   └── fastvideo_fasth3_8step_v2_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`}
+
+

+ Sample with 8 steps and guidance scale 1 (the trained schedule). Video + is fixed 24 fps and frame counts snap down to the 17n+5 grid. +

+
+ ), + }, { name: "ltx2", label: "LTX-2",