diff --git a/README.es.md b/README.es.md
index 4ec1c5d53..1c5d9d4f8 100644
--- a/README.es.md
+++ b/README.es.md
@@ -110,6 +110,7 @@ SimpleTuner es compatible con las siguientes familias de modelos. El soporte det
| **LTX Video 2** | 19B | Apache-2.0 | Sí |
| **Lumina2** | 2B | Apache-2.0 | Sí |
| **Mage-Flow** | 4B | MIT | Sí |
+| **MiniMax H3** | 33B | MiniMax H3 Community License | Aplican condiciones (exclusiones territoriales; autorización requerida en EE. UU./UE/Reino Unido/Corea del Sur) |
| **OmniGen** | 3.8B | MIT | Sí |
| **PixArt Sigma** | 0.6B-0.9B | OpenRAIL++ | Sí (restringido) |
| **Qwen Image** | 20B | Apache-2.0 | Sí |
diff --git a/README.hi.md b/README.hi.md
index 699d5bb07..929abbfd4 100644
--- a/README.hi.md
+++ b/README.hi.md
@@ -110,6 +110,7 @@ SimpleTuner निम्नलिखित मॉडल families का सम
| **LTX Video 2** | 19B | Apache-2.0 | हाँ |
| **Lumina2** | 2B | Apache-2.0 | हाँ |
| **Mage-Flow** | 4B | MIT | हाँ |
+| **MiniMax H3** | 33B | MiniMax H3 Community License | शर्तें लागू (territory exclusions; US/EU/UK/KR में authorization आवश्यक) |
| **OmniGen** | 3.8B | MIT | हाँ |
| **PixArt Sigma** | 0.6B-0.9B | OpenRAIL++ | हाँ (restricted) |
| **Qwen Image** | 20B | Apache-2.0 | हाँ |
diff --git a/README.ja.md b/README.ja.md
index c68739a3a..f7fe8a2de 100644
--- a/README.ja.md
+++ b/README.ja.md
@@ -110,6 +110,7 @@ SimpleTunerは以下のモデルファミリーをサポートしています。
| **LTX Video 2** | 19B | Apache-2.0 | 可 |
| **Lumina2** | 2B | Apache-2.0 | 可 |
| **Mage-Flow** | 4B | MIT | 可 |
+| **MiniMax H3** | 33B | MiniMax H3 Community License | 条件付き（地域除外あり；米国/EU/英国/韓国は認可が必要） |
| **OmniGen** | 3.8B | MIT | 可 |
| **PixArt Sigma** | 0.6B-0.9B | OpenRAIL++ | 可（制限あり） |
| **Qwen Image** | 20B | Apache-2.0 | 可 |
diff --git a/README.md b/README.md
index 3a0e88972..2a7049f6e 100644
--- a/README.md
+++ b/README.md
@@ -110,6 +110,7 @@ SimpleTuner supports the following model families. Detailed training feature sup
| **LTX Video 2** | 19B | Apache-2.0 | Yes |
| **Lumina2** | 2B | Apache-2.0 | Yes |
| **Mage-Flow** | 4B | MIT | Yes |
+| **MiniMax H3** | 33B | MiniMax H3 Community License | Conditions apply (territory exclusions; authorization required in US/EU/UK/KR) |
| **OmniGen** | 3.8B | MIT | Yes |
| **PixArt Sigma** | 0.6B-0.9B | OpenRAIL++ | Yes (restricted) |
| **Qwen Image** | 20B | Apache-2.0 | Yes |
diff --git a/README.pt-BR.md b/README.pt-BR.md
index a4ccf22aa..d5ca9093b 100644
--- a/README.pt-BR.md
+++ b/README.pt-BR.md
@@ -110,6 +110,7 @@ SimpleTuner suporta as seguintes familias de modelos. O suporte detalhado a recu
| **LTX Video 2** | 19B | Apache-2.0 | Sim |
| **Lumina2** | 2B | Apache-2.0 | Sim |
| **Mage-Flow** | 4B | MIT | Sim |
+| **MiniMax H3** | 33B | MiniMax H3 Community License | Condicoes aplicaveis (exclusoes territoriais; autorizacao exigida nos EUA/UE/Reino Unido/Coreia do Sul) |
| **OmniGen** | 3.8B | MIT | Sim |
| **PixArt Sigma** | 0.6B-0.9B | OpenRAIL++ | Sim (restrito) |
| **Qwen Image** | 20B | Apache-2.0 | Sim |
diff --git a/README.zh.md b/README.zh.md
index e1ce18596..eac51b1bf 100644
--- a/README.zh.md
+++ b/README.zh.md
@@ -110,6 +110,7 @@ SimpleTuner 支持以下模型系列。详细的训练功能支持请参阅[快
| **LTX Video 2** | 19B | Apache-2.0 | 是 |
| **Lumina2** | 2B | Apache-2.0 | 是 |
| **Mage-Flow** | 4B | MIT | 是 |
+| **MiniMax H3** | 33B | MiniMax H3 Community License | 有条件（含地区排除；美国/EU/英国/韩国需要授权） |
| **OmniGen** | 3.8B | MIT | 是 |
| **PixArt Sigma** | 0.6B-0.9B | OpenRAIL++ | 是（受限） |
| **Qwen Image** | 20B | Apache-2.0 | 是 |
diff --git a/documentation/DATALOADER.md b/documentation/DATALOADER.md
index 7de8cb662..7a2da658e 100644
--- a/documentation/DATALOADER.md
+++ b/documentation/DATALOADER.md
@@ -682,8 +682,8 @@ For example, with 4 GPUs, `train_batch_size=4`, and `gradient_accumulation_steps
To automatically adjust `repeats` when your dataset is smaller than the effective batch size, use the `--allow_dataset_oversubscription` flag (documented in [OPTIONS.md](OPTIONS.md#--allow_dataset_oversubscription)).
When enabled, SimpleTuner will:
-- Calculate the minimum repeats needed for training
-- Automatically increase `repeats` to meet the requirement
+- Calculate the minimum repeats needed for each undersized aspect bucket
+- Pad only those buckets to meet the effective batch size
- Log a warning showing the adjustment
- **Respect manually-set repeats values** - if you explicitly configure `repeats` in your dataset config, the automatic adjustment will be skipped
@@ -1370,7 +1370,7 @@ Webshart datasets load WebDataset-style tar shards through the `webshart` packag
- `metadata` is optional and points to a separate metadata location when captions or shard metadata are stored outside the shard source. For Hugging Face metadata repos such as `webshart/conceptual-captions-12m-webdataset-metadata`, pass the repo id; Webshart follows the source shard subfolder layout such as `data/`.
- `metadata_backend` must be `webshart`; it reads dimensions and captions from Webshart metadata.
- `caption_strategy` should be `webshart` to train from metadata captions, or `instanceprompt` to ignore stored captions.
-- `webshart.cache_dir` stores SimpleTuner metadata plus Webshart metadata and shard caches. `shard_cache_gb` and `parallel_downloads` are passed to Webshart's shard cache.
+- `webshart.cache_dir` stores SimpleTuner metadata plus Webshart metadata and shard caches. `shard_cache_gb` and `parallel_downloads` are passed to Webshart's shard cache; set `shard_cache_gb` to `0` to disable whole-shard caching and retain indexed range reads.
This backend requires a Webshart build with `TarDataLoader.list_shard_sample_aspect_buckets()`.
diff --git a/documentation/OPTIONS.es.md b/documentation/OPTIONS.es.md
index 61bb37a60..826a9e4c1 100644
--- a/documentation/OPTIONS.es.md
+++ b/documentation/OPTIONS.es.md
@@ -52,6 +52,15 @@ Donde `foo` es tu entorno de configuración; o simplemente usa `config/config.js
- `diffusers` es el esquema estándar de PEFT/Diffusers.
- `comfyui` convierte hacia/desde claves estilo ComfyUI (`diffusion_model.*` con tensores `lora_A/lora_B` y `.alpha`). Flux, Flux2, Lumina2 y Z-Image detectarán automáticamente entradas ComfyUI incluso si esto se deja en `diffusers`, pero cámbialo a `comfyui` para forzar salida ComfyUI al guardar.
+### `--minimax_h3_target_mode`
+
+- **Qué**: Controla si MiniMax-H3 incluye filas de audio objetivo.
+- **Opciones**: `auto`, `video`, `av`
+- **Predeterminado**: `auto`
+- **Notas**:
+ - `auto` se resuelve como solo video, omitiendo caché VAE de audio, colación y filas de audio objetivo para H3.
+ - Define `minimax_h3_target_mode` o `h3_target_mode` como `av` en una entrada de data backend para activar entrenamiento conjunto audio-video en un backend de audio auto-split o explícito.
+
### `--fuse_qkv_projections`
- **Qué**: Fusiona las proyecciones QKV en los bloques de atención del modelo para un uso más eficiente del hardware.
@@ -1880,6 +1889,7 @@ usage: train.py [-h] --model_family
[--flow_beta_schedule_beta FLOW_BETA_SCHEDULE_BETA]
[--flow_schedule_shift FLOW_SCHEDULE_SHIFT]
[--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]]
+ [--audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT]
[--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS]
[--flow_timesteps_mode {fixed-list,round-robin}]
[--flux_guidance_mode {constant,random-range}]
@@ -2000,7 +2010,7 @@ usage: train.py [-h] --model_family
[--rescale_betas_zero_snr [RESCALE_BETAS_ZERO_SNR]]
[--webhook_config WEBHOOK_CONFIG]
[--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL]
- [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}]
+ [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}]
[--distillation_config DISTILLATION_CONFIG]
[--ema_validation {none,ema_only,comparison}]
[--local_rank LOCAL_RANK] [--ltx_train_mode {t2v,i2v}]
@@ -2351,6 +2361,9 @@ options:
Shift the noise schedule for flow-matching models
--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]
Auto-adjust schedule shift based on image resolution
+ --audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT
+ Shift the audio noise schedule for flow-matching
+ models with audio latents
--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS
Override flow-matching timestep sampling with a fixed
comma-separated list. The list is interpreted as
@@ -2733,7 +2746,7 @@ options:
Path to webhook configuration file
--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL
Interval for webhook reports (seconds)
- --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}
+ --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}
Method for model distillation
Distillation methods cannot be combined with
--train_text_encoder.
diff --git a/documentation/OPTIONS.hi.md b/documentation/OPTIONS.hi.md
index c9377a630..a3fbbd6db 100644
--- a/documentation/OPTIONS.hi.md
+++ b/documentation/OPTIONS.hi.md
@@ -52,6 +52,15 @@ simpletuner configure config/foo/config.json
- `diffusers` standard PEFT/Diffusers layout है।
- `comfyui` keys को ComfyUI‑style में convert करता है (`diffusion_model.*` के साथ `lora_A/lora_B` और `.alpha` tensors)। Flux, Flux2, Lumina2, और Z‑Image ComfyUI inputs को auto‑detect करेंगे भले ही यह `diffusers` पर हो, लेकिन saving के लिए ComfyUI output force करने के लिए `comfyui` सेट करें।
+### `--minimax_h3_target_mode`
+
+- **What**: MiniMax-H3 target audio rows शामिल करे या नहीं, इसे नियंत्रित करता है।
+- **Choices**: `auto`, `video`, `av`
+- **Default**: `auto`
+- **Notes**:
+ - `auto` video-only में resolve होता है, जिससे H3 के लिए audio VAE cache, collate, और target audio rows skip होते हैं।
+ - auto-split या explicit audio backend को joint audio-video training में opt in करने के लिए data backend entry में `minimax_h3_target_mode` या `h3_target_mode` को `av` सेट करें।
+
### `--fuse_qkv_projections`
- **What**: मॉडल के attention blocks में QKV projections को fuse करता है ताकि hardware का अधिक कुशल उपयोग हो।
@@ -1878,6 +1887,7 @@ usage: train.py [-h] --model_family
[--flow_beta_schedule_beta FLOW_BETA_SCHEDULE_BETA]
[--flow_schedule_shift FLOW_SCHEDULE_SHIFT]
[--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]]
+ [--audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT]
[--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS]
[--flow_timesteps_mode {fixed-list,round-robin}]
[--flux_guidance_mode {constant,random-range}]
@@ -1998,7 +2008,7 @@ usage: train.py [-h] --model_family
[--rescale_betas_zero_snr [RESCALE_BETAS_ZERO_SNR]]
[--webhook_config WEBHOOK_CONFIG]
[--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL]
- [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}]
+ [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}]
[--distillation_config DISTILLATION_CONFIG]
[--ema_validation {none,ema_only,comparison}]
[--local_rank LOCAL_RANK] [--ltx_train_mode {t2v,i2v}]
@@ -2349,6 +2359,9 @@ options:
Shift the noise schedule for flow-matching models
--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]
Auto-adjust schedule shift based on image resolution
+ --audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT
+ Shift the audio noise schedule for flow-matching
+ models with audio latents
--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS
Override flow-matching timestep sampling with a fixed
comma-separated list. The list is interpreted as
@@ -2731,7 +2744,7 @@ options:
Path to webhook configuration file
--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL
Interval for webhook reports (seconds)
- --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}
+ --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}
Method for model distillation
Distillation methods cannot be combined with
--train_text_encoder.
diff --git a/documentation/OPTIONS.ja.md b/documentation/OPTIONS.ja.md
index d3f435367..66f2987d0 100644
--- a/documentation/OPTIONS.ja.md
+++ b/documentation/OPTIONS.ja.md
@@ -52,6 +52,15 @@ simpletuner configure config/foo/config.json
- `diffusers` は標準の PEFT/Diffusers 形式です。
- `comfyui` は ComfyUI 形式（`diffusion_model.*` と `lora_A/lora_B` + `.alpha`）に変換します。Flux、Flux2、Lumina2、Z-Image は `diffusers` のままでも ComfyUI 入力を自動検出しますが、保存時に ComfyUI 出力を強制したい場合は `comfyui` を指定してください。
+### `--minimax_h3_target_mode`
+
+- **内容**: MiniMax-H3 がターゲット音声行を含めるかを制御します。
+- **選択肢**: `auto`, `video`, `av`
+- **既定**: `auto`
+- **注記**:
+ - `auto` は video-only として扱われ、H3 の audio VAE cache、collate、ターゲット音声行を省略します。
+ - auto-split または明示的な audio backend で joint audio-video training を使う場合は、data backend entry に `minimax_h3_target_mode` または `h3_target_mode` を `av` として設定します。
+
### `--fuse_qkv_projections`
- **内容**: モデルのアテンションブロック内 QKV 投影を融合し、ハードウェア効率を高めます。
@@ -1881,6 +1890,7 @@ usage: train.py [-h] --model_family
[--flow_beta_schedule_beta FLOW_BETA_SCHEDULE_BETA]
[--flow_schedule_shift FLOW_SCHEDULE_SHIFT]
[--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]]
+ [--audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT]
[--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS]
[--flow_timesteps_mode {fixed-list,round-robin}]
[--flux_guidance_mode {constant,random-range}]
@@ -2001,7 +2011,7 @@ usage: train.py [-h] --model_family
[--rescale_betas_zero_snr [RESCALE_BETAS_ZERO_SNR]]
[--webhook_config WEBHOOK_CONFIG]
[--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL]
- [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}]
+ [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}]
[--distillation_config DISTILLATION_CONFIG]
[--ema_validation {none,ema_only,comparison}]
[--local_rank LOCAL_RANK] [--ltx_train_mode {t2v,i2v}]
@@ -2351,6 +2361,9 @@ options:
Shift the noise schedule for flow-matching models
--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]
Auto-adjust schedule shift based on image resolution
+ --audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT
+ Shift the audio noise schedule for flow-matching
+ models with audio latents
--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS
Override flow-matching timestep sampling with a fixed
comma-separated list. The list is interpreted as
@@ -2733,7 +2746,7 @@ options:
Path to webhook configuration file
--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL
Interval for webhook reports (seconds)
- --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}
+ --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}
Method for model distillation
Distillation methods cannot be combined with
--train_text_encoder.
diff --git a/documentation/OPTIONS.md b/documentation/OPTIONS.md
index 71308a0f3..3d61bc2da 100644
--- a/documentation/OPTIONS.md
+++ b/documentation/OPTIONS.md
@@ -52,6 +52,15 @@ Where `foo` is your config environment - or just use `config/config.json` if you
- `diffusers` is the standard PEFT/Diffusers layout.
- `comfyui` converts to/from ComfyUI-style keys (`diffusion_model.*` with `lora_A/lora_B` and `.alpha` tensors). Flux, Flux2, Lumina2, and Z-Image will auto-detect ComfyUI inputs even if this is left at `diffusers`, but set it to `comfyui` to force ComfyUI output when saving.
+### `--minimax_h3_target_mode`
+
+- **What**: Controls whether MiniMax-H3 includes target audio rows.
+- **Choices**: `auto`, `video`, `av`
+- **Default**: `auto`
+- **Notes**:
+ - `auto` resolves to video-only, skipping audio VAE caching, collation, and target audio rows for H3.
+ - Set `minimax_h3_target_mode` or `h3_target_mode` to `av` in a data backend entry to opt an auto-split or explicit audio backend into joint audio-video training.
+
### `--fuse_qkv_projections`
- **What**: Fuses the QKV projections in the model's attention blocks to make more efficient use of hardware.
@@ -1884,6 +1893,7 @@ usage: train.py [-h] --model_family
[--flow_beta_schedule_beta FLOW_BETA_SCHEDULE_BETA]
[--flow_schedule_shift FLOW_SCHEDULE_SHIFT]
[--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]]
+ [--audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT]
[--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS]
[--flow_timesteps_mode {fixed-list,round-robin}]
[--flux_guidance_mode {constant,random-range}]
@@ -2004,7 +2014,7 @@ usage: train.py [-h] --model_family
[--rescale_betas_zero_snr [RESCALE_BETAS_ZERO_SNR]]
[--webhook_config WEBHOOK_CONFIG]
[--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL]
- [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}]
+ [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}]
[--distillation_config DISTILLATION_CONFIG]
[--ema_validation {none,ema_only,comparison}]
[--local_rank LOCAL_RANK] [--ltx_train_mode {t2v,i2v}]
@@ -2355,6 +2365,9 @@ options:
Shift the noise schedule for flow-matching models
--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]
Auto-adjust schedule shift based on image resolution
+ --audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT
+ Shift the audio noise schedule for flow-matching
+ models with audio latents
--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS
Override flow-matching timestep sampling with a fixed
comma-separated list. The list is interpreted as
@@ -2739,7 +2752,7 @@ options:
Path to webhook configuration file
--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL
Interval for webhook reports (seconds)
- --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}
+ --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}
Method for model distillation
Distillation methods cannot be combined with
--train_text_encoder.
diff --git a/documentation/OPTIONS.pt-BR.md b/documentation/OPTIONS.pt-BR.md
index 10c5e132a..44e5c3c43 100644
--- a/documentation/OPTIONS.pt-BR.md
+++ b/documentation/OPTIONS.pt-BR.md
@@ -52,6 +52,15 @@ Onde `foo` e seu ambiente de config — ou use `config/config.json` se nao estiv
- `diffusers` e o layout padrao PEFT/Diffusers.
- `comfyui` converte para/de chaves estilo ComfyUI (`diffusion_model.*` com tensores `lora_A/lora_B` e `.alpha`). Flux, Flux2, Lumina2 e Z-Image auto-detectam entradas ComfyUI mesmo se isso ficar em `diffusers`, mas defina `comfyui` para forcar saida ComfyUI ao salvar.
+### `--minimax_h3_target_mode`
+
+- **O que**: Controla se MiniMax-H3 inclui linhas de audio alvo.
+- **Opcoes**: `auto`, `video`, `av`
+- **Padrao**: `auto`
+- **Notas**:
+ - `auto` resolve para somente video, pulando cache de audio VAE, collate e linhas de audio alvo para H3.
+ - Defina `minimax_h3_target_mode` ou `h3_target_mode` como `av` em uma entrada de data backend para ativar treino conjunto audio-video em um backend de audio auto-split ou explicito.
+
### `--fuse_qkv_projections`
- **O que**: Faz fusao das projecoes QKV nos blocos de atencao do modelo para usar hardware de forma mais eficiente.
@@ -1876,6 +1885,7 @@ usage: train.py [-h] --model_family
[--flow_beta_schedule_beta FLOW_BETA_SCHEDULE_BETA]
[--flow_schedule_shift FLOW_SCHEDULE_SHIFT]
[--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]]
+ [--audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT]
[--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS]
[--flow_timesteps_mode {fixed-list,round-robin}]
[--flux_guidance_mode {constant,random-range}]
@@ -1996,7 +2006,7 @@ usage: train.py [-h] --model_family
[--rescale_betas_zero_snr [RESCALE_BETAS_ZERO_SNR]]
[--webhook_config WEBHOOK_CONFIG]
[--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL]
- [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}]
+ [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}]
[--distillation_config DISTILLATION_CONFIG]
[--ema_validation {none,ema_only,comparison}]
[--local_rank LOCAL_RANK] [--ltx_train_mode {t2v,i2v}]
@@ -2346,6 +2356,9 @@ options:
Shift the noise schedule for flow-matching models
--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]
Auto-adjust schedule shift based on image resolution
+ --audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT
+ Shift the audio noise schedule for flow-matching
+ models with audio latents
--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS
Override flow-matching timestep sampling with a fixed
comma-separated list. The list is interpreted as
@@ -2728,7 +2741,7 @@ options:
Path to webhook configuration file
--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL
Interval for webhook reports (seconds)
- --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}
+ --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}
Method for model distillation
Distillation methods cannot be combined with
--train_text_encoder.
diff --git a/documentation/OPTIONS.zh.md b/documentation/OPTIONS.zh.md
index a72e43fd4..8782a15b1 100644
--- a/documentation/OPTIONS.zh.md
+++ b/documentation/OPTIONS.zh.md
@@ -52,6 +52,15 @@ simpletuner configure config/foo/config.json
- `diffusers` 为标准 PEFT/Diffusers 格式。
- `comfyui` 会转换为 ComfyUI 风格键（`diffusion_model.*`，含 `lora_A/lora_B` 与 `.alpha` 张量）。Flux、Flux2、Lumina2、Z-Image 即便保持 `diffusers` 也会自动识别 ComfyUI 输入，但若希望保存时强制 ComfyUI 输出，请设为 `comfyui`。
+### `--minimax_h3_target_mode`
+
+- **内容**：控制 MiniMax-H3 是否包含目标音频行。
+- **选项**：`auto`, `video`, `av`
+- **默认**：`auto`
+- **说明**：
+ - `auto` 会解析为仅视频，跳过 H3 的音频 VAE 缓存、collate 和目标音频行。
+ - 如需让 auto-split 或显式音频 backend 进行联合音视频训练，请在 data backend 条目中将 `minimax_h3_target_mode` 或 `h3_target_mode` 设为 `av`。
+
### `--fuse_qkv_projections`
- **内容**：融合注意力块中的 QKV 投影，提高硬件效率。
@@ -1883,6 +1892,7 @@ usage: train.py [-h] --model_family
[--flow_beta_schedule_beta FLOW_BETA_SCHEDULE_BETA]
[--flow_schedule_shift FLOW_SCHEDULE_SHIFT]
[--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]]
+ [--audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT]
[--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS]
[--flow_timesteps_mode {fixed-list,round-robin}]
[--flux_guidance_mode {constant,random-range}]
@@ -2003,7 +2013,7 @@ usage: train.py [-h] --model_family
[--rescale_betas_zero_snr [RESCALE_BETAS_ZERO_SNR]]
[--webhook_config WEBHOOK_CONFIG]
[--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL]
- [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}]
+ [--distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}]
[--distillation_config DISTILLATION_CONFIG]
[--ema_validation {none,ema_only,comparison}]
[--local_rank LOCAL_RANK] [--ltx_train_mode {t2v,i2v}]
@@ -2353,6 +2363,9 @@ options:
Shift the noise schedule for flow-matching models
--flow_schedule_auto_shift [FLOW_SCHEDULE_AUTO_SHIFT]
Auto-adjust schedule shift based on image resolution
+ --audio_flow_schedule_shift AUDIO_FLOW_SCHEDULE_SHIFT
+ Shift the audio noise schedule for flow-matching
+ models with audio latents
--flow_custom_timesteps FLOW_CUSTOM_TIMESTEPS
Override flow-matching timestep sampling with a fixed
comma-separated list. The list is interpreted as
@@ -2735,7 +2748,7 @@ options:
Path to webhook configuration file
--webhook_reporting_interval WEBHOOK_REPORTING_INTERVAL
Interval for webhook reports (seconds)
- --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow}
+ --distillation_method {lcm,dcm,dmd,perflow,flow_dpo,anyflow,h3_drift}
Method for model distillation
Distillation methods cannot be combined with
--train_text_encoder.
diff --git a/documentation/QUICKSTART.es.md b/documentation/QUICKSTART.es.md
index 01870c9e3..ef2a6cd0e 100644
--- a/documentation/QUICKSTART.es.md
+++ b/documentation/QUICKSTART.es.md
@@ -31,6 +31,7 @@
| LTX Video 2 | 19B | [LTXVIDEO2.es.md](quickstart/LTXVIDEO2.es.md) |
| Lumina2 | 2B | [LUMINA2.es.md](quickstart/LUMINA2.es.md) |
| Mage-Flow | 4B | [MAGEFLOW.es.md](quickstart/MAGEFLOW.es.md) |
+| MiniMax H3 | 33B | [MINIMAX_H3.es.md](/documentation/quickstart/MINIMAX_H3.es.md) |
| OmniGen | 3.8B | [OMNIGEN.es.md](quickstart/OMNIGEN.es.md) |
| PixArt Sigma | 0.6B-0.9B | [SIGMA.es.md](quickstart/SIGMA.es.md) |
| Qwen Image | 20B | [QWEN_IMAGE.es.md](quickstart/QWEN_IMAGE.es.md)
[QWEN_EDIT.es.md](quickstart/QWEN_EDIT.es.md) |
@@ -80,6 +81,7 @@ La matriz completa de compatibilidad se divide por área de función para que ca
| LTX Video 2 | ✓ | ✓ | ✓* | ✗ | ✓ opt |
| Lumina2 | ✓ | ✓ | ✓ | ✗ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓* | ✗ | ✓ edit |
+| MiniMax H3 | ✓ | ✓ | ✓* | ✗ | ✓ opt (FL2VA/Ref2VA) |
| OmniGen | ✓ | ✓ | ✓ | ✗ | ✗ |
| PixArt Sigma | ✗ | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓* | ✗ | ✓ req (Edit) |
@@ -127,6 +129,7 @@ La matriz completa de compatibilidad se divide por área de función para que ca
| LTX Video 2 | int8/fp8 optional | bf16 |
| Lumina2 | int8 optional | bf16 |
| Mage-Flow | fp8 optional | bf16 |
+| MiniMax H3 | int8/fp8 optional; convrot-int8 | bf16 |
| OmniGen | int8/fp8 optional | bf16 |
| PixArt Sigma | int8 optional | bf16 |
| Qwen Image | required (int8/nf4) | bf16 |
@@ -174,6 +177,7 @@ La matriz completa de compatibilidad se divide por área de función para que ca
| LTX Video 2 | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | ✓ | ✓ | ✓ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓ | ✓ |
+| MiniMax H3 | ✓ | ✓ | ✓ | ✓ |
| OmniGen | ✓ | ✗ | ✗ | ✗ |
| PixArt Sigma | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓ | ✗ |
@@ -221,6 +225,7 @@ La matriz completa de compatibilidad se divide por área de función para que ca
| LTX Video 2 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| Mage-Flow | flow matching | ✓ | ✓ | ✗ | ✓ | ✓ |
+| MiniMax H3 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| OmniGen | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| PixArt Sigma | epsilon | ✗ | ✗ | ✓ | ✓ | ✓ |
| Qwen Image | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -268,6 +273,7 @@ La matriz completa de compatibilidad se divide por área de función para que ca
| LTX Video 2 | Gemma3 | not specified | LTX Video 2 VAE |
| Lumina2 | Gemma2 | 2B | AutoencoderKL |
| Mage-Flow | Qwen3-VL | not specified | Mage-VAE |
+| MiniMax H3 | Qwen3-VL | not specified | MiniMax H3 Video VAE + Audio VAE |
| OmniGen | Integrated OmniGen encoder | not specified | AutoencoderKL |
| PixArt Sigma | T5 XXL v1.1 | 11B | AutoencoderKL |
| Qwen Image | Qwen2.5-VL | 7B | Qwen Image VAE |
diff --git a/documentation/QUICKSTART.hi.md b/documentation/QUICKSTART.hi.md
index e404d878c..5cc568eaa 100644
--- a/documentation/QUICKSTART.hi.md
+++ b/documentation/QUICKSTART.hi.md
@@ -31,6 +31,7 @@
| LTX Video 2 | 19B | [LTXVIDEO2.hi.md](quickstart/LTXVIDEO2.hi.md) |
| Lumina2 | 2B | [LUMINA2.hi.md](quickstart/LUMINA2.hi.md) |
| Mage-Flow | 4B | [MAGEFLOW.hi.md](quickstart/MAGEFLOW.hi.md) |
+| MiniMax H3 | 33B | [MINIMAX_H3.hi.md](/documentation/quickstart/MINIMAX_H3.hi.md) |
| OmniGen | 3.8B | [OMNIGEN.hi.md](quickstart/OMNIGEN.hi.md) |
| PixArt Sigma | 0.6B-0.9B | [SIGMA.hi.md](quickstart/SIGMA.hi.md) |
| Qwen Image | 20B | [QWEN_IMAGE.hi.md](quickstart/QWEN_IMAGE.hi.md)
[QWEN_EDIT.hi.md](quickstart/QWEN_EDIT.hi.md) |
@@ -80,6 +81,7 @@
| LTX Video 2 | ✓ | ✓ | ✓* | ✗ | ✓ opt |
| Lumina2 | ✓ | ✓ | ✓ | ✗ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓* | ✗ | ✓ edit |
+| MiniMax H3 | ✓ | ✓ | ✓* | ✗ | ✓ opt (FL2VA/Ref2VA) |
| OmniGen | ✓ | ✓ | ✓ | ✗ | ✗ |
| PixArt Sigma | ✗ | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓* | ✗ | ✓ req (Edit) |
@@ -127,6 +129,7 @@
| LTX Video 2 | int8/fp8 optional | bf16 |
| Lumina2 | int8 optional | bf16 |
| Mage-Flow | fp8 optional | bf16 |
+| MiniMax H3 | int8/fp8 optional; convrot-int8 | bf16 |
| OmniGen | int8/fp8 optional | bf16 |
| PixArt Sigma | int8 optional | bf16 |
| Qwen Image | required (int8/nf4) | bf16 |
@@ -174,6 +177,7 @@
| LTX Video 2 | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | ✓ | ✓ | ✓ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓ | ✓ |
+| MiniMax H3 | ✓ | ✓ | ✓ | ✓ |
| OmniGen | ✓ | ✗ | ✗ | ✗ |
| PixArt Sigma | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓ | ✗ |
@@ -221,6 +225,7 @@
| LTX Video 2 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| Mage-Flow | flow matching | ✓ | ✓ | ✗ | ✓ | ✓ |
+| MiniMax H3 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| OmniGen | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| PixArt Sigma | epsilon | ✗ | ✗ | ✓ | ✓ | ✓ |
| Qwen Image | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -268,6 +273,7 @@
| LTX Video 2 | Gemma3 | not specified | LTX Video 2 VAE |
| Lumina2 | Gemma2 | 2B | AutoencoderKL |
| Mage-Flow | Qwen3-VL | not specified | Mage-VAE |
+| MiniMax H3 | Qwen3-VL | not specified | MiniMax H3 Video VAE + Audio VAE |
| OmniGen | Integrated OmniGen encoder | not specified | AutoencoderKL |
| PixArt Sigma | T5 XXL v1.1 | 11B | AutoencoderKL |
| Qwen Image | Qwen2.5-VL | 7B | Qwen Image VAE |
diff --git a/documentation/QUICKSTART.ja.md b/documentation/QUICKSTART.ja.md
index e26c33129..08b9d8beb 100644
--- a/documentation/QUICKSTART.ja.md
+++ b/documentation/QUICKSTART.ja.md
@@ -31,6 +31,7 @@
| LTX Video 2 | 19B | [LTXVIDEO2.ja.md](quickstart/LTXVIDEO2.ja.md) |
| Lumina2 | 2B | [LUMINA2.ja.md](quickstart/LUMINA2.ja.md) |
| Mage-Flow | 4B | [MAGEFLOW.ja.md](quickstart/MAGEFLOW.ja.md) |
+| MiniMax H3 | 33B | [MINIMAX_H3.ja.md](/documentation/quickstart/MINIMAX_H3.ja.md) |
| OmniGen | 3.8B | [OMNIGEN.ja.md](quickstart/OMNIGEN.ja.md) |
| PixArt Sigma | 0.6B-0.9B | [SIGMA.ja.md](quickstart/SIGMA.ja.md) |
| Qwen Image | 20B | [QWEN_IMAGE.ja.md](quickstart/QWEN_IMAGE.ja.md)
[QWEN_EDIT.ja.md](quickstart/QWEN_EDIT.ja.md) |
@@ -80,6 +81,7 @@
| LTX Video 2 | ✓ | ✓ | ✓* | ✗ | ✓ opt |
| Lumina2 | ✓ | ✓ | ✓ | ✗ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓* | ✗ | ✓ edit |
+| MiniMax H3 | ✓ | ✓ | ✓* | ✗ | ✓ opt (FL2VA/Ref2VA) |
| OmniGen | ✓ | ✓ | ✓ | ✗ | ✗ |
| PixArt Sigma | ✗ | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓* | ✗ | ✓ req (Edit) |
@@ -127,6 +129,7 @@
| LTX Video 2 | int8/fp8 optional | bf16 |
| Lumina2 | int8 optional | bf16 |
| Mage-Flow | fp8 optional | bf16 |
+| MiniMax H3 | int8/fp8 optional; convrot-int8 | bf16 |
| OmniGen | int8/fp8 optional | bf16 |
| PixArt Sigma | int8 optional | bf16 |
| Qwen Image | required (int8/nf4) | bf16 |
@@ -174,6 +177,7 @@
| LTX Video 2 | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | ✓ | ✓ | ✓ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓ | ✓ |
+| MiniMax H3 | ✓ | ✓ | ✓ | ✓ |
| OmniGen | ✓ | ✗ | ✗ | ✗ |
| PixArt Sigma | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓ | ✗ |
@@ -221,6 +225,7 @@
| LTX Video 2 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| Mage-Flow | flow matching | ✓ | ✓ | ✗ | ✓ | ✓ |
+| MiniMax H3 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| OmniGen | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| PixArt Sigma | epsilon | ✗ | ✗ | ✓ | ✓ | ✓ |
| Qwen Image | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -268,6 +273,7 @@
| LTX Video 2 | Gemma3 | not specified | LTX Video 2 VAE |
| Lumina2 | Gemma2 | 2B | AutoencoderKL |
| Mage-Flow | Qwen3-VL | not specified | Mage-VAE |
+| MiniMax H3 | Qwen3-VL | not specified | MiniMax H3 Video VAE + Audio VAE |
| OmniGen | Integrated OmniGen encoder | not specified | AutoencoderKL |
| PixArt Sigma | T5 XXL v1.1 | 11B | AutoencoderKL |
| Qwen Image | Qwen2.5-VL | 7B | Qwen Image VAE |
diff --git a/documentation/QUICKSTART.md b/documentation/QUICKSTART.md
index 240dacf64..199c82321 100644
--- a/documentation/QUICKSTART.md
+++ b/documentation/QUICKSTART.md
@@ -31,6 +31,7 @@
| LTX Video 2 | 19B | [LTXVIDEO2.md](/documentation/quickstart/LTXVIDEO2.md) |
| Lumina2 | 2B | [LUMINA2.md](/documentation/quickstart/LUMINA2.md) |
| Mage-Flow | 4B | [MAGEFLOW.md](/documentation/quickstart/MAGEFLOW.md) |
+| MiniMax H3 | 33B | [MINIMAX_H3.md](/documentation/quickstart/MINIMAX_H3.md) |
| OmniGen | 3.8B | [OMNIGEN.md](/documentation/quickstart/OMNIGEN.md) |
| PixArt Sigma | 0.6B-0.9B | [SIGMA.md](/documentation/quickstart/SIGMA.md) |
| Qwen Image | 20B | [QWEN_IMAGE.md](/documentation/quickstart/QWEN_IMAGE.md)
[QWEN_EDIT.md](/documentation/quickstart/QWEN_EDIT.md) |
@@ -80,6 +81,7 @@ The complete compatibility matrix is split by feature area so each table stays r
| LTX Video 2 | ✓ | ✓ | ✓* | ✗ | ✓ opt |
| Lumina2 | ✓ | ✓ | ✓ | ✗ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓* | ✗ | ✓ edit |
+| MiniMax H3 | ✓ | ✓ | ✓* | ✗ | ✓ opt (FL2VA/Ref2VA) |
| OmniGen | ✓ | ✓ | ✓ | ✗ | ✗ |
| PixArt Sigma | ✗ | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓* | ✗ | ✓ req (Edit) |
@@ -127,6 +129,7 @@ The complete compatibility matrix is split by feature area so each table stays r
| LTX Video 2 | int8/fp8 optional | bf16 |
| Lumina2 | int8 optional | bf16 |
| Mage-Flow | fp8 optional | bf16 |
+| MiniMax H3 | int8/fp8 optional; convrot-int8 | bf16 |
| OmniGen | int8/fp8 optional | bf16 |
| PixArt Sigma | int8 optional | bf16 |
| Qwen Image | required (int8/nf4) | bf16 |
@@ -174,6 +177,7 @@ The complete compatibility matrix is split by feature area so each table stays r
| LTX Video 2 | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | ✓ | ✓ | ✓ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓ | ✓ |
+| MiniMax H3 | ✓ | ✓ | ✓ | ✓ |
| OmniGen | ✓ | ✗ | ✗ | ✗ |
| PixArt Sigma | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓ | ✗ |
@@ -221,6 +225,7 @@ The complete compatibility matrix is split by feature area so each table stays r
| LTX Video 2 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| Mage-Flow | flow matching | ✓ | ✓ | ✗ | ✓ | ✓ |
+| MiniMax H3 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| OmniGen | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| PixArt Sigma | epsilon | ✗ | ✗ | ✓ | ✓ | ✓ |
| Qwen Image | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -268,6 +273,7 @@ The complete compatibility matrix is split by feature area so each table stays r
| LTX Video 2 | Gemma3 | not specified | LTX Video 2 VAE |
| Lumina2 | Gemma2 | 2B | AutoencoderKL |
| Mage-Flow | Qwen3-VL | not specified | Mage-VAE |
+| MiniMax H3 | Qwen3-VL | not specified | MiniMax H3 Video VAE + Audio VAE |
| OmniGen | Integrated OmniGen encoder | not specified | AutoencoderKL |
| PixArt Sigma | T5 XXL v1.1 | 11B | AutoencoderKL |
| Qwen Image | Qwen2.5-VL | 7B | Qwen Image VAE |
diff --git a/documentation/QUICKSTART.pt-BR.md b/documentation/QUICKSTART.pt-BR.md
index 370599c88..6d62f2821 100644
--- a/documentation/QUICKSTART.pt-BR.md
+++ b/documentation/QUICKSTART.pt-BR.md
@@ -31,6 +31,7 @@
| LTX Video 2 | 19B | [LTXVIDEO2.pt-BR.md](quickstart/LTXVIDEO2.pt-BR.md) |
| Lumina2 | 2B | [LUMINA2.pt-BR.md](quickstart/LUMINA2.pt-BR.md) |
| Mage-Flow | 4B | [MAGEFLOW.pt-BR.md](quickstart/MAGEFLOW.pt-BR.md) |
+| MiniMax H3 | 33B | [MINIMAX_H3.pt-BR.md](/documentation/quickstart/MINIMAX_H3.pt-BR.md) |
| OmniGen | 3.8B | [OMNIGEN.pt-BR.md](quickstart/OMNIGEN.pt-BR.md) |
| PixArt Sigma | 0.6B-0.9B | [SIGMA.pt-BR.md](quickstart/SIGMA.pt-BR.md) |
| Qwen Image | 20B | [QWEN_IMAGE.pt-BR.md](quickstart/QWEN_IMAGE.pt-BR.md)
[QWEN_EDIT.pt-BR.md](quickstart/QWEN_EDIT.pt-BR.md) |
@@ -80,6 +81,7 @@ A matriz completa de compatibilidade é dividida por área de recurso para mante
| LTX Video 2 | ✓ | ✓ | ✓* | ✗ | ✓ opt |
| Lumina2 | ✓ | ✓ | ✓ | ✗ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓* | ✗ | ✓ edit |
+| MiniMax H3 | ✓ | ✓ | ✓* | ✗ | ✓ opt (FL2VA/Ref2VA) |
| OmniGen | ✓ | ✓ | ✓ | ✗ | ✗ |
| PixArt Sigma | ✗ | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓* | ✗ | ✓ req (Edit) |
@@ -127,6 +129,7 @@ A matriz completa de compatibilidade é dividida por área de recurso para mante
| LTX Video 2 | int8/fp8 optional | bf16 |
| Lumina2 | int8 optional | bf16 |
| Mage-Flow | fp8 optional | bf16 |
+| MiniMax H3 | int8/fp8 optional; convrot-int8 | bf16 |
| OmniGen | int8/fp8 optional | bf16 |
| PixArt Sigma | int8 optional | bf16 |
| Qwen Image | required (int8/nf4) | bf16 |
@@ -174,6 +177,7 @@ A matriz completa de compatibilidade é dividida por área de recurso para mante
| LTX Video 2 | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | ✓ | ✓ | ✓ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓ | ✓ |
+| MiniMax H3 | ✓ | ✓ | ✓ | ✓ |
| OmniGen | ✓ | ✗ | ✗ | ✗ |
| PixArt Sigma | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓ | ✗ |
@@ -221,6 +225,7 @@ A matriz completa de compatibilidade é dividida por área de recurso para mante
| LTX Video 2 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| Mage-Flow | flow matching | ✓ | ✓ | ✗ | ✓ | ✓ |
+| MiniMax H3 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| OmniGen | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| PixArt Sigma | epsilon | ✗ | ✗ | ✓ | ✓ | ✓ |
| Qwen Image | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -268,6 +273,7 @@ A matriz completa de compatibilidade é dividida por área de recurso para mante
| LTX Video 2 | Gemma3 | not specified | LTX Video 2 VAE |
| Lumina2 | Gemma2 | 2B | AutoencoderKL |
| Mage-Flow | Qwen3-VL | not specified | Mage-VAE |
+| MiniMax H3 | Qwen3-VL | not specified | MiniMax H3 Video VAE + Audio VAE |
| OmniGen | Integrated OmniGen encoder | not specified | AutoencoderKL |
| PixArt Sigma | T5 XXL v1.1 | 11B | AutoencoderKL |
| Qwen Image | Qwen2.5-VL | 7B | Qwen Image VAE |
diff --git a/documentation/QUICKSTART.zh.md b/documentation/QUICKSTART.zh.md
index 47c3366f7..ed0b4c3d3 100644
--- a/documentation/QUICKSTART.zh.md
+++ b/documentation/QUICKSTART.zh.md
@@ -31,6 +31,7 @@
| LTX Video 2 | 19B | [LTXVIDEO2.zh.md](quickstart/LTXVIDEO2.zh.md) |
| Lumina2 | 2B | [LUMINA2.zh.md](quickstart/LUMINA2.zh.md) |
| Mage-Flow | 4B | [MAGEFLOW.zh.md](quickstart/MAGEFLOW.zh.md) |
+| MiniMax H3 | 33B | [MINIMAX_H3.zh.md](/documentation/quickstart/MINIMAX_H3.zh.md) |
| OmniGen | 3.8B | [OMNIGEN.zh.md](quickstart/OMNIGEN.zh.md) |
| PixArt Sigma | 0.6B-0.9B | [SIGMA.zh.md](quickstart/SIGMA.zh.md) |
| Qwen Image | 20B | [QWEN_IMAGE.zh.md](quickstart/QWEN_IMAGE.zh.md)
[QWEN_EDIT.zh.md](quickstart/QWEN_EDIT.zh.md) |
@@ -80,6 +81,7 @@
| LTX Video 2 | ✓ | ✓ | ✓* | ✗ | ✓ opt |
| Lumina2 | ✓ | ✓ | ✓ | ✗ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓* | ✗ | ✓ edit |
+| MiniMax H3 | ✓ | ✓ | ✓* | ✗ | ✓ opt (FL2VA/Ref2VA) |
| OmniGen | ✓ | ✓ | ✓ | ✗ | ✗ |
| PixArt Sigma | ✗ | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓* | ✗ | ✓ req (Edit) |
@@ -127,6 +129,7 @@
| LTX Video 2 | int8/fp8 optional | bf16 |
| Lumina2 | int8 optional | bf16 |
| Mage-Flow | fp8 optional | bf16 |
+| MiniMax H3 | int8/fp8 optional; convrot-int8 | bf16 |
| OmniGen | int8/fp8 optional | bf16 |
| PixArt Sigma | int8 optional | bf16 |
| Qwen Image | required (int8/nf4) | bf16 |
@@ -174,6 +177,7 @@
| LTX Video 2 | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | ✓ | ✓ | ✓ | ✗ |
| Mage-Flow | ✓ | ✓ | ✓ | ✓ |
+| MiniMax H3 | ✓ | ✓ | ✓ | ✓ |
| OmniGen | ✓ | ✗ | ✗ | ✗ |
| PixArt Sigma | ✓ | ✓ | ✓ | ✗ |
| Qwen Image | ✓ | ✓ | ✓ | ✗ |
@@ -221,6 +225,7 @@
| LTX Video 2 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| Lumina2 | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| Mage-Flow | flow matching | ✓ | ✓ | ✗ | ✓ | ✓ |
+| MiniMax H3 | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
| OmniGen | flow matching | ✓ | ✗ | ✓ | ✓ | ✓ |
| PixArt Sigma | epsilon | ✗ | ✗ | ✓ | ✓ | ✓ |
| Qwen Image | flow matching | ✓ | ✓ | ✓ | ✓ | ✓ |
@@ -268,6 +273,7 @@
| LTX Video 2 | Gemma3 | not specified | LTX Video 2 VAE |
| Lumina2 | Gemma2 | 2B | AutoencoderKL |
| Mage-Flow | Qwen3-VL | not specified | Mage-VAE |
+| MiniMax H3 | Qwen3-VL | not specified | MiniMax H3 Video VAE + Audio VAE |
| OmniGen | Integrated OmniGen encoder | not specified | AutoencoderKL |
| PixArt Sigma | T5 XXL v1.1 | 11B | AutoencoderKL |
| Qwen Image | Qwen2.5-VL | 7B | Qwen Image VAE |
diff --git a/documentation/distillation/MINIMAX_H3_DRIFT.es.md b/documentation/distillation/MINIMAX_H3_DRIFT.es.md
new file mode 100644
index 000000000..38b8ffff4
--- /dev/null
+++ b/documentation/distillation/MINIMAX_H3_DRIFT.es.md
@@ -0,0 +1,110 @@
+# Destilación de deriva de MiniMax H3
+
+MiniMax H3 es un modelo de video/audio de flow matching ya destilado. En un entrenamiento LoRA o LyCORIS normal, el adaptador aprende el objetivo del dataset, pero puede mover demasiado el comportamiento destilado del checkpoint base: guidance, equilibrio entre modalidades y la secuencia empaquetada de video/audio.
+
+`h3_drift` evita esa deriva comparando la predicción del adaptador con la predicción del mismo modelo cuando el adaptador está desactivado. No carga otro teacher ni usa una caché de destilación. En cada batch SimpleTuner:
+
+1. calcula la pérdida normal de MiniMax H3 con el adaptador activo;
+2. desactiva temporalmente el adaptador;
+3. ejecuta la base congelada con `torch.no_grad()` y el mismo batch preparado;
+4. calcula MSE entre las predicciones de video/audio;
+5. reactiva el adaptador y retropropaga la pérdida combinada.
+
+```text
+total = sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+Con un distiller interno:
+
+```text
+total = inner_distiller_loss + sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+## Cuándo usarlo
+
+Úsalo para LoRA o LyCORIS de MiniMax H3 salvo que quieras quitar o reemplazar la destilación original. Es útil para LoRAs de estilo o concepto, FL2VA/Ref2VA, entrenamiento conjunto audio/video y flavours cuantizados como `convrot-int8` o `convrot-int4`.
+
+No se admite full-rank: si el transformer completo se actualiza, ya no existe una ruta base congelada fiable para comparar.
+
+## Configuración rápida
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+Los ejemplos incluidos lo activan por defecto con `loss_weight: 0.5`. Es un punto medio: el objetivo del dataset sigue siendo principal, pero la referencia base tiene peso suficiente para limitar la deriva.
+
+## Claves
+
+- `loss_weight`: peso de la pérdida contra la base congelada. Empieza con `0.25` a `0.5`; usa `1.0` si la validación pierde el comportamiento base.
+- `sft_loss_weight`: peso de la pérdida normal. Normalmente debe quedar en `1.0`.
+- `balance`: `token` promedia por elementos válidos; `modality` promedia por modalidad después de aplicar pesos.
+- `video_weight`: peso de la deriva de video.
+- `audio_weight`: peso de la deriva de audio.
+- `inner_distillation_method`: distiller opcional que se ejecuta dentro de `h3_drift`, por ejemplo `anyflow`, `dmd`, `perflow`, `flow_dpo` o `self_forcing`.
+- `inner_distillation_config`: configuración que se pasa al distiller interno.
+
+## Componer otro distiller
+
+`h3_drift` puede envolver otro distiller para usar step distillation o un objetivo de preferencia sin dejar de conservar el comportamiento base de MiniMax H3:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "inner_distillation_method": "anyflow",
+ "inner_distillation_config": {
+ "target_mode": "linear",
+ "r_timestep_sampler": "zero",
+ "loss_weight": 1.0
+ }
+ }
+ }
+}
+```
+
+El wrapper delega preparación de batches, scheduler de validación, caché de distillation, batches de captions y hooks de ciclo de vida al distiller interno. El distiller interno mantiene sus propias comprobaciones de compatibilidad.
+
+`sft_loss_weight` sigue siendo el objetivo normal de MiniMax H3 incluso si el distiller interno reescribe `target`. Si el distiller interno agrega conditioning de timestep FlowMap/AnyFlow, H3 drift calcula ese término SFT con otro forward del adaptador, quitando antes las claves internas de timestep. Así la ruta normal de inferencia H3 a 30 steps sigue anclada mientras se entrena la ruta step-distilled.
+
+## Audio y video
+
+`minimax_h3_target_mode: "auto"` se resuelve como video-only. Usa `"video"` para desactivar audio o `"av"` para entrenar filas objetivo de audio junto con video. También puedes definir `h3_target_mode` o `minimax_h3_target_mode` por data backend.
+
+El distiller sigue el batch preparado: compara solo video en batches video-only, compara video y audio en batches `av`, respeta `audio_latent_mask`, `sample_weight` y las máscaras visuales.
+
+## Mantener la destilación CFG
+
+MiniMax H3 es CFG-distilled. El checkpoint base se valida normalmente con `validation_guidance: 1.0`, `validation_guidance_real: 1.0` y `validation_disable_unconditional: true`. Negative prompting no forma parte del contrato base.
+
+SimpleTuner soporta CFG real y negative prompt para checkpoints que la comunidad pueda reentrenar fuera de esa destilación. `h3_drift` empuja en sentido contrario: mantiene la predicción condicional cerca de la base. Si quieres de-destilar H3 o enseñar negative prompts, reduce `loss_weight` o desactiva este distiller.
+
+## Logs y coste
+
+Los logs importantes son `h3_drift_loss`, `h3_drift_video_loss`, `h3_drift_audio_loss`, los contadores de elementos, `h3_drift_weighted_loss`, `h3_drift_sft_loss`, `h3_drift_inner_total` cuando hay distiller interno, y `total`.
+
+El coste es una pasada forward adicional por step, sin cargar un segundo transformer. Si envuelve un distiller FlowMap/AnyFlow interno con `sft_loss_weight` activo, también ejecuta un forward normal del adaptador para el ancla SFT. Sigue siendo compatible con ConvRot, RamTorch, musubi block swap, gradient checkpointing y attention offload, pero cada preset debe medirse porque los forwards extra pueden cambiar el backend más rápido.
+
+## Problemas comunes
+
+- Error de low-rank: usa `model_type: "lora"`.
+- Audio loss cero: el batch es video-only, el target mode no es `av`, o `audio_latent_mask` excluye todo.
+- El adaptador aprende poco: baja `loss_weight`, sube rank o entrena más.
+- El audio deriva: prueba `balance: "modality"` o sube `audio_weight`.
diff --git a/documentation/distillation/MINIMAX_H3_DRIFT.hi.md b/documentation/distillation/MINIMAX_H3_DRIFT.hi.md
new file mode 100644
index 000000000..85fd4f80b
--- /dev/null
+++ b/documentation/distillation/MINIMAX_H3_DRIFT.hi.md
@@ -0,0 +1,110 @@
+# MiniMax H3 Drift Distillation
+
+MiniMax H3 एक distilled flow-matching video/audio model है। सामान्य LoRA या LyCORIS training में adapter dataset target सीखता है, लेकिन base checkpoint का distilled behavior drift कर सकता है: guidance behavior, modality balance, और packed video/audio sequence layout बदल सकते हैं।
+
+`h3_drift` adapter prediction को उसी model की frozen-base prediction से compare करता है जब adapter disabled होता है। यह अलग teacher checkpoint load नहीं करता और distillation cache इस्तेमाल नहीं करता। हर batch में SimpleTuner:
+
+1. adapter enabled path से normal MiniMax H3 SFT loss निकालता है;
+2. adapter अस्थायी रूप से disable करता है;
+3. उसी prepared batch को `torch.no_grad()` में frozen base से चलाता है;
+4. video/audio predictions के बीच MSE निकालता है;
+5. adapter फिर enable करके combined loss backpropagate करता है।
+
+```text
+total = sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+Inner distiller enabled हो तो:
+
+```text
+total = inner_distiller_loss + sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+## कब इस्तेमाल करें
+
+MiniMax H3 LoRA या LyCORIS training में इसे default रखें, जब तक आपका लक्ष्य original distillation को हटाना न हो। यह style/concept LoRAs, FL2VA/Ref2VA, joint audio/video training, और `convrot-int8` / `convrot-int4` जैसे quantized flavours के लिए उपयोगी है।
+
+Full-rank H3 drift supported नहीं है। जब पूरा transformer train हो रहा हो, frozen base comparison path भरोसेमंद नहीं रहता।
+
+## Quick Config
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+Checked-in H3 examples इसे default रूप से enable करते हैं। `loss_weight: 0.5` normal dataset target को primary रखता है, लेकिन base-reference drift को भी meaningful बनाता है।
+
+## Config Keys
+
+- `loss_weight`: frozen-base prediction loss का multiplier। Narrow LoRA के लिए `0.25` से `0.5` शुरू करें; base behavior टूटे तो `1.0` करें।
+- `sft_loss_weight`: normal MiniMax H3 training loss का multiplier। सामान्य fine-tuning में `1.0` रखें।
+- `balance`: `token` valid elements से average करता है; `modality` video और audio means को modality level पर balance करता है।
+- `video_weight`: video drift term का multiplier।
+- `audio_weight`: audio drift term का multiplier।
+- `inner_distillation_method`: optional distiller जो `h3_drift` के अंदर चलेगा, जैसे `anyflow`, `dmd`, `perflow`, `flow_dpo`, या `self_forcing`।
+- `inner_distillation_config`: inner distiller को दी जाने वाली config।
+
+## दूसरा Distiller Compose करना
+
+`h3_drift` किसी और distiller को wrap कर सकता है, ताकि step distillation या preference objective के साथ MiniMax H3 का frozen-base behavior भी preserve रहे।
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "inner_distillation_method": "anyflow",
+ "inner_distillation_config": {
+ "target_mode": "linear",
+ "r_timestep_sampler": "zero",
+ "loss_weight": 1.0
+ }
+ }
+ }
+}
+```
+
+Wrapper batch preparation, validation scheduler, distillation cache, caption batches, और generator/discriminator lifecycle hooks inner distiller को delegate करता है। Inner distiller की compatibility checks फिर भी लागू रहती हैं।
+
+Inner distiller `target` rewrite करे तब भी `sft_loss_weight` normal MiniMax H3 objective ही रहता है। अगर inner distiller FlowMap/AnyFlow timestep conditioning जोड़ता है, तो H3 drift inner timestep keys हटाकर adapter-enabled forward फिर चलाता है और SFT term निकालता है। इससे step-distilled path train होता है, लेकिन normal 30-step H3 inference path anchored रहता है।
+
+## Video और Audio Modes
+
+`minimax_h3_target_mode: "auto"` video-only बनता है। `"video"` audio target rows बंद रखता है। `"av"` joint audio/video rows train करता है। इसे global config में या data backend में `h3_target_mode` / `minimax_h3_target_mode` से set कर सकते हैं।
+
+Distiller prepared batch follow करता है: video-only batch में सिर्फ `model_prediction`, `av` batch में video और `audio_prediction`, `audio_latent_mask`, `sample_weight`, और visual masks सभी लागू होते हैं।
+
+## CFG Distillation बनाए रखना
+
+MiniMax H3 CFG-distilled है। Base checkpoint आम तौर पर `validation_guidance: 1.0`, `validation_guidance_real: 1.0`, और `validation_disable_unconditional: true` के साथ validate होता है। Negative prompting base contract का हिस्सा नहीं है।
+
+SimpleTuner real CFG और negative prompt encode कर सकता है क्योंकि community H3 को de-distill कर सकती है। `h3_drift` उल्टा pressure देता है: adapter को base conditional prediction के पास रखता है। अगर आपका लक्ष्य negative prompt behavior सिखाना या de-distillation है, तो `loss_weight` घटाएं या distiller disable करें।
+
+## Logs और Cost
+
+मुख्य logs: `h3_drift_loss`, `h3_drift_video_loss`, `h3_drift_audio_loss`, element counts, `h3_drift_weighted_loss`, `h3_drift_sft_loss`, inner distiller enabled होने पर `h3_drift_inner_total`, और `total`।
+
+हर step में एक extra forward pass लगता है, लेकिन दूसरा transformer memory में नहीं रखा जाता। FlowMap/AnyFlow inner distiller को `sft_loss_weight` enabled के साथ wrap करने पर SFT anchor के लिए normal adapter forward भी चलता है। ConvRot, RamTorch, musubi block swap, gradient checkpointing, और attention offload के साथ यह compatible है; फिर भी presets benchmark करें क्योंकि extra forwards fastest backend बदल सकते हैं।
+
+## Troubleshooting
+
+- Low-rank error: `model_type: "lora"` इस्तेमाल करें।
+- Audio loss zero: batch video-only है, target mode `av` नहीं है, या `audio_latent_mask` सब exclude कर रहा है।
+- Adapter concept कम सीख रहा है: `loss_weight` घटाएं, rank बढ़ाएं, या training लंबी करें।
+- Audio drift कर रहा है: `balance: "modality"` या बड़ा `audio_weight` आजमाएं।
diff --git a/documentation/distillation/MINIMAX_H3_DRIFT.ja.md b/documentation/distillation/MINIMAX_H3_DRIFT.ja.md
new file mode 100644
index 000000000..0d767d4b9
--- /dev/null
+++ b/documentation/distillation/MINIMAX_H3_DRIFT.ja.md
@@ -0,0 +1,110 @@
+# MiniMax H3 Drift Distillation
+
+MiniMax H3 は distilled flow-matching の video/audio モデルです。通常の LoRA / LyCORIS training では adapter が dataset target を学習しますが、base checkpoint の distilled behavior が drift することがあります。guidance behavior、modality balance、packed video/audio sequence layout が必要以上に変わる場合があります。
+
+`h3_drift` は、adapter 有効時の prediction を、adapter を無効化した同じ model の frozen-base prediction と比較します。別の teacher checkpoint は読み込まず、distillation cache も使いません。各 batch で SimpleTuner は次を行います。
+
+1. adapter 有効の通常 path で MiniMax H3 SFT loss を計算する。
+2. adapter を一時的に無効化する。
+3. 同じ prepared batch を `torch.no_grad()` の frozen base で実行する。
+4. video/audio prediction の MSE を計算する。
+5. adapter を戻し、combined loss を backpropagate する。
+
+```text
+total = sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+inner distiller を使う場合:
+
+```text
+total = inner_distiller_loss + sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+## 使う場面
+
+MiniMax H3 の LoRA / LyCORIS training では、元の distillation behavior を意図的に外す場合を除き有効化してください。style/concept LoRA、FL2VA/Ref2VA、joint audio/video training、`convrot-int8` / `convrot-int4` などの quantized flavour に向いています。
+
+Full-rank では未対応です。transformer 全体を更新する場合、比較対象となる frozen base path が成立しません。
+
+## Quick Config
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+同梱 H3 examples はこれを default で有効化します。`loss_weight: 0.5` は dataset target を主にしつつ、base reference drift を抑える実用的な開始値です。
+
+## Config Keys
+
+- `loss_weight`: frozen-base prediction loss の倍率。狭い LoRA は `0.25` から `0.5`、base behavior が崩れる場合は `1.0`。
+- `sft_loss_weight`: 通常の MiniMax H3 training loss の倍率。通常は `1.0`。
+- `balance`: `token` は valid element 数で平均、`modality` は video/audio の modality mean を重み付きで平均。
+- `video_weight`: video drift term の倍率。
+- `audio_weight`: audio drift term の倍率。
+- `inner_distillation_method`: `h3_drift` の内側で実行する optional distiller。例: `anyflow`、`dmd`、`perflow`、`flow_dpo`、`self_forcing`。
+- `inner_distillation_config`: inner distiller に渡す config。
+
+## 他の Distiller との合成
+
+`h3_drift` は、MiniMax H3 の frozen-base behavior を保ちながら step distillation や preference objective を使うために、別の distiller を wrap できます。
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "inner_distillation_method": "anyflow",
+ "inner_distillation_config": {
+ "target_mode": "linear",
+ "r_timestep_sampler": "zero",
+ "loss_weight": 1.0
+ }
+ }
+ }
+}
+```
+
+Wrapper は batch preparation、validation scheduler、distillation cache、caption batch support、generator/discriminator lifecycle hooks を inner distiller に委譲します。inner distiller の compatibility check はそのまま有効です。
+
+inner distiller が `target` を書き換えても、`sft_loss_weight` は通常の MiniMax H3 objective のままです。inner distiller が FlowMap/AnyFlow timestep conditioning を追加する場合、H3 drift は inner timestep key を外したうえで adapter 有効の forward をもう一度実行し、SFT term を計算します。これにより step-distilled path を学習しながら、標準の 30-step H3 inference path を維持できます。
+
+## Video と Audio
+
+`minimax_h3_target_mode: "auto"` は video-only に解決されます。`"video"` は audio target rows を使わず、`"av"` は joint audio/video rows を学習します。global config または data backend の `h3_target_mode` / `minimax_h3_target_mode` で指定できます。
+
+Distiller は prepared batch に従います。video-only batch では `model_prediction` のみ、`av` batch では video と `audio_prediction` を比較します。`audio_latent_mask`、`sample_weight`、visual mask も反映されます。
+
+## CFG Distillation を保つ
+
+MiniMax H3 は CFG-distilled です。base checkpoint は通常 `validation_guidance: 1.0`、`validation_guidance_real: 1.0`、`validation_disable_unconditional: true` で検証します。Negative prompting は base contract に含まれません。
+
+SimpleTuner は real CFG と negative prompt encoding に対応しています。community が H3 を de-distill する可能性があるためです。`h3_drift` は逆方向の制約で、adapter を base conditional prediction に近づけます。negative prompt behavior を教えたい、または de-distillation したい場合は `loss_weight` を下げるか無効化してください。
+
+## Logs と Cost
+
+主な logs は `h3_drift_loss`、`h3_drift_video_loss`、`h3_drift_audio_loss`、element counts、`h3_drift_weighted_loss`、`h3_drift_sft_loss`、inner distiller 有効時の `h3_drift_inner_total`、`total` です。
+
+各 step に extra forward pass が 1 回追加されますが、2 つ目の transformer は保持しません。FlowMap/AnyFlow inner distiller を `sft_loss_weight` 有効で wrap する場合、SFT anchor 用の通常 adapter forward も追加されます。ConvRot、RamTorch、musubi block swap、gradient checkpointing、attention offload と併用できます。ただし extra forward により fastest backend が変わるため、preset ごとに benchmark してください。
+
+## Troubleshooting
+
+- low-rank error: `model_type: "lora"` を使います。
+- audio loss が zero: batch が video-only、target mode が `av` ではない、または `audio_latent_mask` が全行を除外しています。
+- adapter が concept を学びにくい: `loss_weight` を下げる、rank を上げる、または longer training。
+- audio が drift する: `balance: "modality"` または高い `audio_weight` を試します。
diff --git a/documentation/distillation/MINIMAX_H3_DRIFT.md b/documentation/distillation/MINIMAX_H3_DRIFT.md
new file mode 100644
index 000000000..d77185cd9
--- /dev/null
+++ b/documentation/distillation/MINIMAX_H3_DRIFT.md
@@ -0,0 +1,155 @@
+# MiniMax H3 Drift Distillation
+
+MiniMax H3 is a distilled flow-matching video/audio model. During normal LoRA or LyCORIS training, the adapter learns the dataset target directly, but the base model's distilled behavior can drift: guidance behavior, modality balance, and the packed video/audio sequence layout can move away from the frozen checkpoint even when the training loss looks reasonable.
+
+H3 drift distillation keeps the adapter close to the frozen MiniMax H3 prediction while still allowing the normal supervised fine-tuning loss to teach the dataset. It is not a separate student/teacher checkpoint and it does not cache teacher latents. SimpleTuner uses the same model twice on the same prepared batch:
+
+1. Run the trainable adapter path and compute the normal MiniMax H3 SFT loss.
+2. Temporarily disable the adapter.
+3. Run the frozen base path under `torch.no_grad()` with the same prompt embeddings, noised latents, timesteps, masks, and packed layout.
+4. Compare the adapter prediction to the frozen-base prediction with a video/audio MSE loss.
+5. Re-enable the adapter and backpropagate the combined loss.
+
+The combined objective is:
+
+```text
+total = sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+When an inner distiller is configured, the objective becomes:
+
+```text
+total = inner_distiller_loss + sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+This makes the adapter learn the dataset while being penalized for changing unrelated base behavior too aggressively.
+
+## When To Use It
+
+Use `h3_drift` for MiniMax H3 LoRA or LyCORIS training unless you are intentionally trying to remove or replace the original distillation behavior. It is especially useful for:
+
+- short concept or style LoRAs where the adapter should not change the model's guidance behavior;
+- FL2VA or Ref2VA training where the model must keep its packed conditioning layout stable;
+- joint audio/video training where video can otherwise dominate the loss simply because it has more elements;
+- quantized H3 flavours such as `convrot-int8` and `convrot-int4`, where adapter-only training is the expected path.
+
+Do not use it for full-rank MiniMax H3 training. The implementation intentionally rejects non-LoRA training because there is no frozen base path to compare against once the full transformer weights are being updated.
+
+## Quick Config
+
+Add the distiller to a MiniMax H3 LoRA config:
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+The checked-in MiniMax H3 example configs enable this by default. `loss_weight: 0.5` is a practical starter value: the normal dataset loss remains primary, but drift has enough weight to catch adapters that quickly move away from the base prediction.
+
+## Configuration Keys
+
+- `loss_weight`: multiplier for the frozen-base prediction loss. Start at `0.25` to `0.5` for narrow visual LoRAs. Use `1.0` when validation shows the adapter is breaking the base model behavior or when training longer multimodal adapters.
+- `sft_loss_weight`: multiplier for the normal MiniMax H3 training loss. Keep this at `1.0` for ordinary fine-tuning. Setting it to `0.0` turns the run into pure base-prediction following and is usually only useful for debugging.
+- `balance`: `token` or `modality`. `token` averages by valid element count, so video naturally dominates video/audio batches. `modality` averages the video and audio modality means after applying `video_weight` and `audio_weight`.
+- `video_weight`: multiplier for the video prediction drift term.
+- `audio_weight`: multiplier for the audio prediction drift term.
+- `inner_distillation_method`: optional nested SimpleTuner distiller to run inside H3 drift, such as `anyflow`, `dmd`, `perflow`, `flow_dpo`, or `self_forcing`.
+- `inner_distillation_config`: configuration mapping passed to the nested distiller.
+
+The defaults are `loss_weight: 1.0`, `sft_loss_weight: 1.0`, `balance: "token"`, and equal video/audio weights. The examples set a lighter drift value because most downstream LoRAs still need the dataset target to move the adapter.
+
+## Composing Another Distiller
+
+`h3_drift` can wrap another distiller when you want a step-distillation or preference objective while still preserving MiniMax H3's frozen-base behavior. For example, this runs AnyFlow target preparation first, then adds the H3 drift reference loss:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "inner_distillation_method": "anyflow",
+ "inner_distillation_config": {
+ "target_mode": "linear",
+ "r_timestep_sampler": "zero",
+ "loss_weight": 1.0
+ }
+ }
+ }
+}
+```
+
+The wrapper delegates batch preparation, validation scheduler hooks, distillation cache generation, caption-batch support, and discriminator/generator lifecycle hooks to the inner distiller. The inner distiller still performs its own compatibility checks, so a method that does not support MiniMax H3 will fail during setup instead of silently falling back.
+
+`sft_loss_weight` remains the normal MiniMax H3 objective even when the inner distiller rewrites `target`. If the inner distiller adds FlowMap/AnyFlow timestep conditioning, H3 drift computes that SFT term with an additional adapter-enabled forward pass after removing the inner timestep keys. That keeps the standard 30-step H3 inference path anchored while the inner distiller trains the step-distilled path.
+
+## Video And Audio Modes
+
+MiniMax H3 can train video-only targets or joint audio/video targets:
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+`auto` resolves to video-only. Set `minimax_h3_target_mode` globally, or `h3_target_mode` / `minimax_h3_target_mode` inside a data backend entry, to `av` when you want target audio rows included.
+
+The drift distiller follows the prepared batch:
+
+- video-only batches compare only `model_prediction`;
+- audio/video batches compare both `model_prediction` and `audio_prediction`;
+- `audio_latent_mask` masks out missing generated audio rows;
+- existing sample weights and visual loss masks are respected.
+
+If an `av` backend has no audio latents, MiniMax H3 builds zero audio rows and masks audio loss. That avoids training against fake audio targets while still allowing video training to proceed.
+
+## Maintaining CFG-Distilled Behavior
+
+MiniMax H3 is CFG-distilled. The base checkpoint is expected to work with `validation_guidance: 1.0`, `validation_guidance_real: 1.0`, and `validation_disable_unconditional: true`. Negative prompting is not part of the base training contract.
+
+SimpleTuner still has pipeline support for real CFG and negative prompt encoding because users may train adapters or merged checkpoints that partially remove the original distillation. `h3_drift` is the opposite pressure: it helps keep the adapter close to the base model's distilled conditional prediction. If your goal is to teach negative prompt behavior or de-distill H3, reduce `loss_weight` or disable `h3_drift`; otherwise the reference loss will push the adapter back toward the base model.
+
+## Reading The Logs
+
+The distiller adds separate log values:
+
+- `h3_drift_loss`: unweighted combined video/audio drift loss;
+- `h3_drift_video_loss`: video drift mean;
+- `h3_drift_audio_loss`: audio drift mean;
+- `h3_drift_video_elements` and `h3_drift_audio_elements`: valid element counts after masks;
+- `h3_drift_weighted_loss`: `h3_drift_loss * loss_weight`;
+- `h3_drift_sft_loss`: normal MiniMax H3 loss after `sft_loss_weight`;
+- `h3_drift_inner_total`: the nested distiller's `total` value, when `inner_distillation_method` is enabled;
+- `total`: final loss returned to the trainer.
+
+Use these values to diagnose whether the adapter is learning the dataset or mostly fighting the reference. If `h3_drift_loss` climbs while validation quality gets more erratic, increase `loss_weight` or reduce learning rate. If the adapter barely learns the concept, reduce `loss_weight`, raise rank, or train longer.
+
+## Cost And Memory
+
+H3 drift distillation adds one extra forward pass per training step, but it does not keep a second transformer in memory. The reference pass reuses the same model with adapters disabled and no gradients. When wrapping an inner FlowMap/AnyFlow distiller with `sft_loss_weight` enabled, it also runs a normal-path adapter forward for the SFT anchor. Peak VRAM still increases because activations, compilation caches, and temporary buffers can overlap with the normal step, but the cost is closer to extra forwards than a second trainable model.
+
+This makes it compatible with the usual MiniMax H3 memory features: quantized ConvRot checkpoints, RamTorch, musubi block swap, gradient checkpointing, and attention offload. Benchmark each preset after enabling the distiller because the extra forward can change which checkpointing or attention backend is fastest.
+
+## Troubleshooting
+
+- **`H3 drift distillation only supports low-rank LoRA/LyCORIS training`**: set `model_type` to `lora`. Full-rank H3 drift is intentionally unsupported.
+- **`H3 drift prediction contains no target modality`**: the model output did not contain a video or audio prediction. Check `minimax_h3_target_mode`, latent cache shape, and model family.
+- **Audio loss is always zero**: the batch is video-only, `minimax_h3_target_mode` resolves to `video`, or `audio_latent_mask` excludes every audio row.
+- **Adapter learning is too weak**: lower `loss_weight` from `0.5` to `0.25`, increase rank, or use `balance: "token"` if `modality` was over-emphasizing audio.
+- **Audio quality drifts while video looks stable**: use `balance: "modality"` or raise `audio_weight` for joint audio/video training.
diff --git a/documentation/distillation/MINIMAX_H3_DRIFT.pt-BR.md b/documentation/distillation/MINIMAX_H3_DRIFT.pt-BR.md
new file mode 100644
index 000000000..14b020d98
--- /dev/null
+++ b/documentation/distillation/MINIMAX_H3_DRIFT.pt-BR.md
@@ -0,0 +1,110 @@
+# Destilação de Drift do MiniMax H3
+
+MiniMax H3 é um modelo de video/audio flow-matching já destilado. Em training LoRA ou LyCORIS normal, o adapter aprende o target do dataset, mas pode deslocar demais o comportamento destilado do checkpoint base: guidance, balanço de modalidades e o layout da sequência video/audio empacotada.
+
+`h3_drift` compara a prediction do adapter com a prediction do mesmo modelo quando o adapter está desativado. Ele não carrega outro teacher e não usa cache de distillation. A cada batch o SimpleTuner:
+
+1. calcula a loss normal do MiniMax H3 com o adapter ativo;
+2. desativa temporariamente o adapter;
+3. roda a base congelada com `torch.no_grad()` no mesmo prepared batch;
+4. calcula MSE entre predictions de video/audio;
+5. reativa o adapter e faz backprop da loss combinada.
+
+```text
+total = sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+Com um distiller interno:
+
+```text
+total = inner_distiller_loss + sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+## Quando usar
+
+Use em LoRA ou LyCORIS de MiniMax H3, a menos que você queira remover a destilação original. É útil para LoRAs de estilo/conceito, FL2VA/Ref2VA, training conjunto audio/video e flavours quantizados como `convrot-int8` e `convrot-int4`.
+
+Full-rank não é suportado. Quando o transformer inteiro é atualizado, não há uma rota base congelada confiável para comparar.
+
+## Config rápido
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+Os exemplos H3 incluídos ativam isso por padrão. `loss_weight: 0.5` mantém o target do dataset como principal, mas dá peso suficiente à referência base para conter drift.
+
+## Chaves
+
+- `loss_weight`: multiplicador da loss contra a base congelada. Comece com `0.25` a `0.5`; use `1.0` se a validação perder comportamento base.
+- `sft_loss_weight`: multiplicador da loss normal. Normalmente fica em `1.0`.
+- `balance`: `token` faz média por elementos válidos; `modality` faz média por modalidade depois dos pesos.
+- `video_weight`: peso do termo de drift de video.
+- `audio_weight`: peso do termo de drift de audio.
+- `inner_distillation_method`: distiller opcional executado dentro de `h3_drift`, por exemplo `anyflow`, `dmd`, `perflow`, `flow_dpo` ou `self_forcing`.
+- `inner_distillation_config`: configuração passada para o distiller interno.
+
+## Compor outro distiller
+
+`h3_drift` pode envolver outro distiller para usar step distillation ou objetivo de preferência enquanto preserva o comportamento base do MiniMax H3:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "inner_distillation_method": "anyflow",
+ "inner_distillation_config": {
+ "target_mode": "linear",
+ "r_timestep_sampler": "zero",
+ "loss_weight": 1.0
+ }
+ }
+ }
+}
+```
+
+O wrapper delega preparação de batches, scheduler de validação, cache de distillation, batches de captions e hooks de ciclo de vida ao distiller interno. O distiller interno mantém suas próprias validações de compatibilidade.
+
+`sft_loss_weight` continua sendo o objetivo normal do MiniMax H3 mesmo quando o distiller interno reescreve `target`. Se o distiller interno adiciona conditioning de timestep FlowMap/AnyFlow, H3 drift calcula esse termo SFT com outro forward do adapter, removendo antes as chaves internas de timestep. Isso mantém a rota normal de inferência H3 em 30 steps ancorada enquanto a rota step-distilled é treinada.
+
+## Video e Audio
+
+`minimax_h3_target_mode: "auto"` vira video-only. Use `"video"` para não treinar audio target rows, ou `"av"` para joint audio/video. Também pode ser definido por data backend com `h3_target_mode` ou `minimax_h3_target_mode`.
+
+O distiller segue o prepared batch: compara só video em video-only, compara video e `audio_prediction` em `av`, e respeita `audio_latent_mask`, `sample_weight` e visual masks.
+
+## Manter a destilação CFG
+
+MiniMax H3 é CFG-distilled. O checkpoint base normalmente valida com `validation_guidance: 1.0`, `validation_guidance_real: 1.0` e `validation_disable_unconditional: true`. Negative prompting não faz parte do contrato base.
+
+SimpleTuner suporta real CFG e negative prompt encoding porque a comunidade pode de-destilar H3. `h3_drift` faz a pressão oposta: mantém o adapter perto da conditional prediction da base. Para ensinar negative prompts ou de-destilar, reduza `loss_weight` ou desative o distiller.
+
+## Logs e custo
+
+Logs principais: `h3_drift_loss`, `h3_drift_video_loss`, `h3_drift_audio_loss`, element counts, `h3_drift_weighted_loss`, `h3_drift_sft_loss`, `h3_drift_inner_total` quando houver distiller interno, e `total`.
+
+Cada step ganha um forward pass extra, mas não mantém um segundo transformer na memória. Ao envolver um distiller FlowMap/AnyFlow interno com `sft_loss_weight` ativo, também executa um forward normal do adapter para a âncora SFT. Funciona com ConvRot, RamTorch, musubi block swap, gradient checkpointing e attention offload; ainda assim, benchmark cada preset porque os forwards extra podem mudar o backend mais rápido.
+
+## Troubleshooting
+
+- Erro de low-rank: use `model_type: "lora"`.
+- Audio loss zero: batch video-only, target mode não é `av`, ou `audio_latent_mask` exclui tudo.
+- Adapter aprende pouco: reduza `loss_weight`, aumente rank ou treine por mais tempo.
+- Audio deriva: use `balance: "modality"` ou aumente `audio_weight`.
diff --git a/documentation/distillation/MINIMAX_H3_DRIFT.zh.md b/documentation/distillation/MINIMAX_H3_DRIFT.zh.md
new file mode 100644
index 000000000..243cc4a96
--- /dev/null
+++ b/documentation/distillation/MINIMAX_H3_DRIFT.zh.md
@@ -0,0 +1,110 @@
+# MiniMax H3 漂移蒸馏
+
+MiniMax H3 是已经蒸馏过的 flow-matching 视频/音频模型。普通 LoRA 或 LyCORIS 训练会让 adapter 学习数据集目标，但也可能让基础 checkpoint 的蒸馏行为发生漂移：guidance 行为、模态平衡以及打包的视频/音频序列布局都会被过度改变。
+
+`h3_drift` 会把 adapter 启用时的 prediction，与同一个模型在 adapter 关闭时的 frozen-base prediction 对齐。它不加载第二个 teacher，也不使用 distillation cache。每个 batch 中 SimpleTuner 会：
+
+1. 用 adapter 启用路径计算正常 MiniMax H3 SFT loss；
+2. 暂时关闭 adapter；
+3. 用相同 prepared batch 在 `torch.no_grad()` 下运行冻结 base；
+4. 对 video/audio prediction 计算 MSE；
+5. 重新启用 adapter，并反传组合 loss。
+
+```text
+total = sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+启用内部 distiller 时：
+
+```text
+total = inner_distiller_loss + sft_loss_weight * normal_h3_loss + loss_weight * frozen_base_prediction_mse
+```
+
+## 何时使用
+
+除非你明确想移除或替换原始蒸馏行为，否则 MiniMax H3 LoRA / LyCORIS 训练建议启用它。它适合 style/concept LoRA、FL2VA/Ref2VA、联合音视频训练，以及 `convrot-int8` / `convrot-int4` 等量化 flavour。
+
+Full-rank 不支持 H3 drift。更新整个 transformer 时，不再存在可比较的冻结 base 路径。
+
+## 快速配置
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+随仓库提供的 H3 examples 默认启用该 distiller。`loss_weight: 0.5` 会让数据集目标保持主导，同时让 base reference 有足够权重来限制漂移。
+
+## 配置项
+
+- `loss_weight`：冻结 base prediction loss 的倍率。窄域 LoRA 可从 `0.25` 到 `0.5` 开始；如果验证显示 base 行为被破坏，可用 `1.0`。
+- `sft_loss_weight`：正常 MiniMax H3 training loss 的倍率。普通 fine-tuning 保持 `1.0`。
+- `balance`：`token` 按有效元素数平均；`modality` 按 video/audio 模态均值加权平均。
+- `video_weight`：video drift term 权重。
+- `audio_weight`：audio drift term 权重。
+- `inner_distillation_method`：可选的内部 distiller，会在 `h3_drift` 内执行，例如 `anyflow`、`dmd`、`perflow`、`flow_dpo` 或 `self_forcing`。
+- `inner_distillation_config`：传给内部 distiller 的配置。
+
+## 组合另一个 distiller
+
+`h3_drift` 可以包裹另一个 distiller，在使用 step distillation 或偏好目标的同时继续保持 MiniMax H3 的冻结 base 行为：
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "inner_distillation_method": "anyflow",
+ "inner_distillation_config": {
+ "target_mode": "linear",
+ "r_timestep_sampler": "zero",
+ "loss_weight": 1.0
+ }
+ }
+ }
+}
+```
+
+该 wrapper 会把 batch preparation、validation scheduler hook、distillation cache、caption batch 支持以及 generator/discriminator 生命周期 hook 委托给内部 distiller。内部 distiller 仍会执行自己的兼容性检查。
+
+即使内部 distiller 重写了 `target`，`sft_loss_weight` 仍表示正常 MiniMax H3 目标。如果内部 distiller 添加 FlowMap/AnyFlow timestep conditioning，H3 drift 会先移除内部 timestep key，再额外运行一次 adapter-enabled forward 来计算 SFT 项。这样 step-distilled 路径可以训练，同时标准 30-step H3 推理路径仍被锚定。
+
+## 视频与音频
+
+`minimax_h3_target_mode: "auto"` 会解析为 video-only。使用 `"video"` 跳过音频目标行，使用 `"av"` 训练联合音视频目标行。也可以在 data backend 中设置 `h3_target_mode` 或 `minimax_h3_target_mode`。
+
+Distiller 跟随 prepared batch：video-only batch 只比较 `model_prediction`；`av` batch 同时比较 video 和 `audio_prediction`；并尊重 `audio_latent_mask`、`sample_weight` 和视觉 mask。
+
+## 保持 CFG 蒸馏行为
+
+MiniMax H3 是 CFG-distilled。基础 checkpoint 通常使用 `validation_guidance: 1.0`、`validation_guidance_real: 1.0`、`validation_disable_unconditional: true` 验证。Negative prompting 不属于基础训练契约。
+
+SimpleTuner 支持 real CFG 和 negative prompt encoding，是因为社区可能会把 H3 重新训练到不再保持原始蒸馏。`h3_drift` 是相反的约束：它让 adapter 靠近 base conditional prediction。如果目标是训练 negative prompt 行为或 de-distill H3，请降低 `loss_weight` 或关闭该 distiller。
+
+## 日志与成本
+
+主要日志包括 `h3_drift_loss`、`h3_drift_video_loss`、`h3_drift_audio_loss`、元素计数、`h3_drift_weighted_loss`、`h3_drift_sft_loss`、启用内部 distiller 时的 `h3_drift_inner_total` 和 `total`。
+
+每个 step 会增加一次 forward pass，但不会在显存中保存第二个 transformer。若包裹 FlowMap/AnyFlow 内部 distiller 且启用 `sft_loss_weight`，还会为 SFT 锚点运行一次正常 adapter forward。它可与 ConvRot、RamTorch、musubi block swap、gradient checkpointing 和 attention offload 一起使用；不过额外 forward 可能改变最快 backend，因此每个 preset 都应重新 benchmark。
+
+## 排错
+
+- low-rank 错误：使用 `model_type: "lora"`。
+- audio loss 始终为零：batch 是 video-only、target mode 不是 `av`，或 `audio_latent_mask` 排除了所有音频行。
+- adapter 学不到概念：降低 `loss_weight`、提高 rank 或延长训练。
+- 音频发生漂移：尝试 `balance: "modality"` 或提高 `audio_weight`。
diff --git a/documentation/quickstart/MINIMAX_H3.es.md b/documentation/quickstart/MINIMAX_H3.es.md
new file mode 100644
index 000000000..5f2671458
--- /dev/null
+++ b/documentation/quickstart/MINIMAX_H3.es.md
@@ -0,0 +1,89 @@
+# Guía rápida de MiniMax H3
+
+MiniMax H3 es un modelo de video/audio flow-matching de 33B. SimpleTuner soporta training de adapters con la familia `minimaxh3`, incluyendo conditioning FL2VA de primer/último frame y flavours ConvRot cuantizados.
+
+## Configs iniciales
+
+Empieza desde uno de estos ejemplos:
+
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora`
+
+Usa el preset más cercano a tu VRAM y ajusta resolución, frames, attention backend y checkpointing después de un smoke test.
+
+## Ajustes principales
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "mixed_precision": "bf16",
+ "base_model_precision": "no_change",
+ "text_encoder_1_precision": "int8-quanto",
+ "flow_schedule_shift": 12.0,
+ "audio_flow_schedule_shift": 3.0,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0
+}
+```
+
+Los ejemplos usan `convrot-int8`. Puedes usar `convrot-int4` con la misma familia si quieres el checkpoint de menor precisión.
+
+## Mantener la destilación
+
+MiniMax H3 es CFG-distilled. El checkpoint base está pensado para funcionar sin rama unconditional, así que los ejemplos validan con guidance `1.0` y `validation_disable_unconditional: true`.
+
+Los adapters pueden alejarse del comportamiento destilado. Por eso los ejemplos activan `h3_drift` por defecto:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+El distiller ejecuta una referencia de la base congelada con el adapter desactivado y penaliza la deriva de la predicción video/audio. Mantenlo activo para LoRAs normales. Baja `loss_weight` si el concepto no aprende; súbelo si la validación pierde el comportamiento base. Más detalles: [MiniMax H3 Drift Distillation](../distillation/MINIMAX_H3_DRIFT.es.md).
+
+Negative prompting no forma parte del contrato base de H3. SimpleTuner mantiene CFG real y negative prompts para checkpoints de-destilados, pero `h3_drift` preserva el comportamiento condicional original.
+
+## Modo de audio
+
+`minimax_h3_target_mode: "auto"` se resuelve como video-only y evita trabajo de audio VAE:
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+Usa `"av"` solo si el dataset tiene latentes de audio y quieres training conjunto audio/video. También puedes configurarlo por data backend con `h3_target_mode` o `minimax_h3_target_mode`.
+
+## Memoria
+
+- Usa el ejemplo 24G con RamTorch si la VRAM es limitada.
+- Prueba `musubi_blocks_to_swap` antes de subir mucho el checkpointing.
+- Mantén `flow_schedule_shift` de video en `12.0` y `audio_flow_schedule_shift` en `3.0`. El helper H3 corrige el default global heredado `3.0` para video porque no coincide con el schedule de MiniMax H3.
+- SimpleTuner fuerza VAE tiling y temporal roll/chunking para el video VAE de H3. La geometría usa el upstream `256` tile size con `64` overlap; poner esas opciones en false se ignora porque el decode sin tiling puede producir cambios fuertes de color y patrones halftone.
+- Benchmarkea `attention_mechanism` en la GPU real.
+- Repite el smoke test si cambias `torch.compile`, porque las cachés pueden aumentar VRAM.
+
+## Ejecutar
+
+```bash
+simpletuner train example=minimaxh3-fl2va-convrot-int8.peft-lora
+```
+
+Haz un smoke test corto y revisa que `h3_drift_loss`, la pérdida normal y las validaciones se comporten de forma coherente.
diff --git a/documentation/quickstart/MINIMAX_H3.hi.md b/documentation/quickstart/MINIMAX_H3.hi.md
new file mode 100644
index 000000000..b7337ff25
--- /dev/null
+++ b/documentation/quickstart/MINIMAX_H3.hi.md
@@ -0,0 +1,89 @@
+# MiniMax H3 Quickstart
+
+MiniMax H3 एक 33B flow-matching video/audio model है। SimpleTuner `minimaxh3` family में adapter training support करता है, जिसमें FL2VA first/last-frame conditioning और quantized ConvRot flavours शामिल हैं।
+
+## Starting Configs
+
+इन examples से शुरू करें:
+
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora`
+
+अपनी VRAM के सबसे करीब preset लें, फिर smoke test के बाद resolution, frames, attention backend, और checkpointing बदलें।
+
+## Core Settings
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "mixed_precision": "bf16",
+ "base_model_precision": "no_change",
+ "text_encoder_1_precision": "int8-quanto",
+ "flow_schedule_shift": 12.0,
+ "audio_flow_schedule_shift": 3.0,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0
+}
+```
+
+Examples `convrot-int8` use करते हैं। कम precision checkpoint चाहिए तो उसी family में `convrot-int4` इस्तेमाल कर सकते हैं।
+
+## Distillation बनाए रखना
+
+MiniMax H3 CFG-distilled है। Base checkpoint unconditional branch के बिना चलने के लिए बना है, इसलिए examples guidance `1.0` और `validation_disable_unconditional: true` रखते हैं।
+
+Adapter training base distilled behavior से drift कर सकती है। इसलिए examples default रूप से `h3_drift` enable करते हैं:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+यह adapter disabled frozen-base reference pass चलाता है और video/audio prediction drift को penalize करता है। Normal H3 LoRA के लिए इसे enabled रखें। Concept नहीं सीख रहा हो तो `loss_weight` घटाएं; validation base behavior खो रही हो तो बढ़ाएं। पूरी explanation: [MiniMax H3 Drift Distillation](../distillation/MINIMAX_H3_DRIFT.hi.md)।
+
+Negative prompting base H3 contract का हिस्सा नहीं है। SimpleTuner de-distilled checkpoints के लिए real CFG और negative prompts support रखता है, लेकिन `h3_drift` original distilled conditional behavior preserve करता है।
+
+## Audio Target Mode
+
+`minimax_h3_target_mode: "auto"` video-only बनता है और audio VAE work बचाता है:
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+`"av"` तभी use करें जब dataset में target audio latents हों और joint audio/video training चाहिए। Per-backend `h3_target_mode` या `minimax_h3_target_mode` भी set कर सकते हैं।
+
+## Memory Knobs
+
+- VRAM tight हो तो 24G RamTorch example इस्तेमाल करें।
+- ज्यादा checkpointing से पहले `musubi_blocks_to_swap` test करें।
+- video `flow_schedule_shift` को `12.0` और `audio_flow_schedule_shift` को `3.0` रखें। H3 helper inherited global video default `3.0` को ठीक करता है क्योंकि वह MiniMax H3 schedule से match नहीं करता।
+- SimpleTuner H3 video VAE के लिए VAE tiling और temporal roll/chunking force करता है। Tiling geometry upstream जैसी है: `256` tile size और `64` overlap। इन options को false करने पर ignore किया जाएगा, क्योंकि untiled decode से severe colour shift और halftone artifacts आ सकते हैं।
+- Target GPU पर `attention_mechanism` benchmark करें।
+- `torch.compile` बदलने पर smoke test दोबारा करें, क्योंकि compile cache VRAM बदल सकता है।
+
+## Run
+
+```bash
+simpletuner train example=minimaxh3-fl2va-convrot-int8.peft-lora
+```
+
+Long run से पहले छोटा smoke test करें और देखें कि `h3_drift_loss`, normal loss, और validation samples साथ-साथ sensible move कर रहे हैं।
diff --git a/documentation/quickstart/MINIMAX_H3.ja.md b/documentation/quickstart/MINIMAX_H3.ja.md
new file mode 100644
index 000000000..21e1fb1d8
--- /dev/null
+++ b/documentation/quickstart/MINIMAX_H3.ja.md
@@ -0,0 +1,89 @@
+# MiniMax H3 Quickstart
+
+MiniMax H3 は 33B の flow-matching video/audio モデルです。SimpleTuner は `minimaxh3` family で adapter training をサポートし、FL2VA first/last-frame conditioning と quantized ConvRot flavours を扱えます。
+
+## Starting Configs
+
+次の examples から始めます。
+
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora`
+
+VRAM に近い preset を選び、smoke test 後に resolution、frame count、attention backend、checkpointing を調整します。
+
+## Core Settings
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "mixed_precision": "bf16",
+ "base_model_precision": "no_change",
+ "text_encoder_1_precision": "int8-quanto",
+ "flow_schedule_shift": 12.0,
+ "audio_flow_schedule_shift": 3.0,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0
+}
+```
+
+Examples は `convrot-int8` を使います。より低い precision の checkpoint を使う場合は同じ family で `convrot-int4` を指定できます。
+
+## Distillation を保つ
+
+MiniMax H3 は CFG-distilled です。base checkpoint は unconditional branch なしで動く想定なので、examples は guidance `1.0` と `validation_disable_unconditional: true` を使います。
+
+Adapter training は distilled behavior から drift することがあります。そのため examples は default で `h3_drift` を有効化します。
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+これは adapter を無効化した frozen-base reference pass を実行し、video/audio prediction drift を penalize します。通常の H3 LoRA では有効のままにしてください。concept を学びにくい場合は `loss_weight` を下げ、validation が base behavior を失う場合は上げます。詳細は [MiniMax H3 Drift Distillation](../distillation/MINIMAX_H3_DRIFT.ja.md) を参照してください。
+
+Negative prompting は base H3 contract の一部ではありません。SimpleTuner は de-distilled checkpoint 向けに real CFG と negative prompts を残していますが、`h3_drift` は original distilled conditional behavior を保つためのものです。
+
+## Audio Target Mode
+
+`minimax_h3_target_mode: "auto"` は video-only になり、audio VAE work を避けます。
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+dataset に target audio latents があり joint audio/video training したい場合だけ `"av"` を使います。data backend ごとに `h3_target_mode` または `minimax_h3_target_mode` でも設定できます。
+
+## Memory Knobs
+
+- VRAM が厳しい場合は 24G RamTorch example を使います。
+- heavy checkpointing の前に `musubi_blocks_to_swap` を試します。
+- video `flow_schedule_shift` は `12.0`、`audio_flow_schedule_shift` は `3.0` のままにします。H3 helper は、MiniMax H3 schedule と一致しない継承された global video default `3.0` を修正します。
+- SimpleTuner は H3 video VAE の tiling と temporal roll/chunking を強制的に有効化します。tiling geometry は upstream と同じ `256` tile size / `64` overlap です。これらを false にしても無視されます。untiled decode は大きな color shift や halftone artifacts を出すことがあります。
+- 実際の GPU で `attention_mechanism` を benchmark します。
+- `torch.compile` を変えたら smoke test をやり直します。compile cache が VRAM を変えるためです。
+
+## Run
+
+```bash
+simpletuner train example=minimaxh3-fl2va-convrot-int8.peft-lora
+```
+
+long run の前に短い smoke test を行い、`h3_drift_loss`、normal loss、validation sample が一貫して動くことを確認します。
diff --git a/documentation/quickstart/MINIMAX_H3.md b/documentation/quickstart/MINIMAX_H3.md
new file mode 100644
index 000000000..25f10db73
--- /dev/null
+++ b/documentation/quickstart/MINIMAX_H3.md
@@ -0,0 +1,91 @@
+# MiniMax H3 Quickstart
+
+MiniMax H3 is a 33B flow-matching video/audio model. SimpleTuner supports adapter training through the MiniMax H3 model family, including FL2VA-style first/last-frame conditioning and the ConvRot quantized flavours.
+
+## Starting Configs
+
+The checked-in examples are the recommended starting points:
+
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora`
+
+Use the closest VRAM preset first, then adjust resolution, frame count, attention backend, and checkpointing after a smoke test.
+
+## Core Settings
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "mixed_precision": "bf16",
+ "base_model_precision": "no_change",
+ "text_encoder_1_precision": "int8-quanto",
+ "flow_schedule_shift": 12.0,
+ "audio_flow_schedule_shift": 3.0,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0
+}
+```
+
+`convrot-int8` is the default example flavour. `convrot-int4` can be used with the same model family when the lower precision checkpoint is desired.
+
+## Maintaining Distillation
+
+MiniMax H3 is CFG-distilled. The base checkpoint is intended to run without an unconditional branch, so the example configs keep validation at guidance `1.0` and set `validation_disable_unconditional: true`.
+
+Adapter training can still drift away from the frozen distilled behavior. The examples enable the H3 drift distiller by default:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+This runs a no-grad frozen-base reference pass with the adapter disabled and penalizes the adapter when its video/audio prediction moves too far from that reference. Keep it enabled for normal H3 LoRAs. Lower `loss_weight` if the adapter is not learning the concept; raise it if validation starts losing the base model's distilled behavior. For the full explanation, see [MiniMax H3 Drift Distillation](../distillation/MINIMAX_H3_DRIFT.md).
+
+Negative prompting is not part of the base H3 contract. SimpleTuner keeps real CFG and negative prompt plumbing available for de-distilled community checkpoints, but `h3_drift` deliberately preserves the original distilled conditional behavior.
+
+## Audio Target Mode
+
+By default, `minimax_h3_target_mode: "auto"` resolves to video-only and avoids audio VAE work:
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+Use `"av"` only when the dataset has target audio latents and you want joint audio/video training. You can set `h3_target_mode` or `minimax_h3_target_mode` inside a data backend entry to opt only selected backends into audio.
+
+## Memory Knobs
+
+- Use the 24G RamTorch example when VRAM is tight.
+- Use `musubi_blocks_to_swap` when block swap is faster than heavier checkpointing on your GPU.
+- Keep video `flow_schedule_shift` at `12.0` and audio `audio_flow_schedule_shift` at `3.0`. The H3 helper corrects the inherited global `3.0` video default because it does not match the MiniMax H3 schedule.
+- SimpleTuner forces H3 video VAE tiling and temporal roll/chunking on. The tiling geometry is the upstream `256` tile size with `64` overlap; setting those options false is ignored because untiled decode can produce severe colour shifts and halftone artifacts.
+- Benchmark `attention_mechanism` values on the target GPU; H3 shapes may prefer a different backend than Wan or LTX Video.
+- Re-test after changing `torch.compile` mode because compile caches can change peak VRAM.
+
+## Run
+
+From the SimpleTuner directory:
+
+```bash
+simpletuner train example=minimaxh3-fl2va-convrot-int8.peft-lora
+```
+
+Run a short smoke test before committing to a long run. Check that `h3_drift_loss`, normal training loss, and validation samples move together instead of one term dominating the run.
diff --git a/documentation/quickstart/MINIMAX_H3.pt-BR.md b/documentation/quickstart/MINIMAX_H3.pt-BR.md
new file mode 100644
index 000000000..832f3e595
--- /dev/null
+++ b/documentation/quickstart/MINIMAX_H3.pt-BR.md
@@ -0,0 +1,89 @@
+# Guia rápido do MiniMax H3
+
+MiniMax H3 é um modelo de video/audio flow-matching de 33B. SimpleTuner suporta adapter training pela family `minimaxh3`, incluindo conditioning FL2VA de primeiro/último frame e flavours ConvRot quantizados.
+
+## Configs iniciais
+
+Comece por um destes examples:
+
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora`
+
+Use o preset mais próximo da sua VRAM e ajuste resolução, frames, attention backend e checkpointing depois de um smoke test.
+
+## Ajustes principais
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "mixed_precision": "bf16",
+ "base_model_precision": "no_change",
+ "text_encoder_1_precision": "int8-quanto",
+ "flow_schedule_shift": 12.0,
+ "audio_flow_schedule_shift": 3.0,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0
+}
+```
+
+Os examples usam `convrot-int8`. Use `convrot-int4` na mesma family se quiser o checkpoint de menor precision.
+
+## Manter a destilação
+
+MiniMax H3 é CFG-distilled. O checkpoint base foi feito para rodar sem branch unconditional, então os examples validam com guidance `1.0` e `validation_disable_unconditional: true`.
+
+Adapter training ainda pode driftar para longe do comportamento destilado. Por isso os examples ativam `h3_drift` por padrão:
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+O distiller roda uma referência da base congelada com o adapter desativado e penaliza drift na prediction de video/audio. Mantenha ativo para LoRAs H3 normais. Reduza `loss_weight` se o concept não aprende; aumente se a validação perde o comportamento base. Veja [MiniMax H3 Drift Distillation](../distillation/MINIMAX_H3_DRIFT.pt-BR.md).
+
+Negative prompting não faz parte do contrato base do H3. SimpleTuner mantém real CFG e negative prompts para checkpoints de-destilados, mas `h3_drift` preserva o comportamento condicional original.
+
+## Modo de audio
+
+`minimax_h3_target_mode: "auto"` vira video-only e evita trabalho de audio VAE:
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+Use `"av"` só quando o dataset tiver target audio latents e você quiser joint audio/video training. Também pode configurar por backend com `h3_target_mode` ou `minimax_h3_target_mode`.
+
+## Memória
+
+- Use o example 24G com RamTorch quando VRAM estiver apertada.
+- Teste `musubi_blocks_to_swap` antes de aumentar muito checkpointing.
+- Mantenha `flow_schedule_shift` de video em `12.0` e `audio_flow_schedule_shift` em `3.0`. O helper H3 corrige o default global herdado `3.0` para video porque ele não corresponde ao schedule do MiniMax H3.
+- SimpleTuner força VAE tiling e temporal roll/chunking para o video VAE do H3. A geometria usa o upstream `256` tile size com `64` overlap; definir essas opções como false é ignorado porque decode sem tiling pode causar fortes mudanças de cor e padrões halftone.
+- Faça benchmark de `attention_mechanism` na GPU real.
+- Refaça o smoke test após mudar `torch.compile`, porque caches podem aumentar VRAM.
+
+## Rodar
+
+```bash
+simpletuner train example=minimaxh3-fl2va-convrot-int8.peft-lora
+```
+
+Faça um smoke test curto antes de uma execução longa e confirme que `h3_drift_loss`, loss normal e samples de validação evoluem de forma coerente.
diff --git a/documentation/quickstart/MINIMAX_H3.zh.md b/documentation/quickstart/MINIMAX_H3.zh.md
new file mode 100644
index 000000000..3791df572
--- /dev/null
+++ b/documentation/quickstart/MINIMAX_H3.zh.md
@@ -0,0 +1,89 @@
+# MiniMax H3 快速开始
+
+MiniMax H3 是 33B flow-matching 视频/音频模型。SimpleTuner 通过 `minimaxh3` model family 支持 adapter 训练，包括 FL2VA 首/末帧 conditioning 和量化 ConvRot flavour。
+
+## 起始配置
+
+建议从这些 examples 开始：
+
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora`
+- `simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora`
+
+先选择最接近显存的 preset，完成 smoke test 后再调整分辨率、帧数、attention backend 和 checkpointing。
+
+## 核心设置
+
+```json
+{
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "mixed_precision": "bf16",
+ "base_model_precision": "no_change",
+ "text_encoder_1_precision": "int8-quanto",
+ "flow_schedule_shift": 12.0,
+ "audio_flow_schedule_shift": 3.0,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0
+}
+```
+
+Examples 默认使用 `convrot-int8`。如果需要更低 precision checkpoint，可以在同一 model family 中使用 `convrot-int4`。
+
+## 维持蒸馏
+
+MiniMax H3 是 CFG-distilled。基础 checkpoint 预期不使用 unconditional branch，因此 examples 使用 guidance `1.0` 并设置 `validation_disable_unconditional: true`。
+
+Adapter 训练仍可能偏离基础蒸馏行为。因此 examples 默认启用 `h3_drift`：
+
+```json
+{
+ "distillation_method": "h3_drift",
+ "distillation_config": {
+ "h3_drift": {
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0
+ }
+ }
+}
+```
+
+该 distiller 会在 adapter 关闭时运行 no-grad frozen-base reference pass，并惩罚 video/audio prediction 的漂移。普通 H3 LoRA 建议保持启用。adapter 学不到概念时降低 `loss_weight`；验证开始丢失基础行为时提高它。完整说明见 [MiniMax H3 Drift Distillation](../distillation/MINIMAX_H3_DRIFT.zh.md)。
+
+Negative prompting 不属于基础 H3 契约。SimpleTuner 为 de-distilled checkpoints 保留 real CFG 和 negative prompt 支持，但 `h3_drift` 会保留原始蒸馏条件行为。
+
+## 音频目标模式
+
+`minimax_h3_target_mode: "auto"` 会解析为 video-only，并避免 audio VAE 工作：
+
+```json
+{
+ "minimax_h3_target_mode": "video"
+}
+```
+
+只有当 dataset 有 target audio latents 并需要联合音视频训练时才使用 `"av"`。也可以在 data backend 中设置 `h3_target_mode` 或 `minimax_h3_target_mode`。
+
+## 显存选项
+
+- 显存紧张时使用 24G RamTorch example。
+- 在加重 checkpointing 前测试 `musubi_blocks_to_swap`。
+- video `flow_schedule_shift` 保持 `12.0`，`audio_flow_schedule_shift` 保持 `3.0`。H3 helper 会修正继承来的全局 video 默认值 `3.0`，因为它不匹配 MiniMax H3 schedule。
+- SimpleTuner 会强制启用 H3 video VAE tiling 和 temporal roll/chunking。tiling 几何与 upstream 一致，使用 `256` tile size 和 `64` overlap；把这些选项设为 false 会被忽略，因为未 tiling 的 decode 可能产生严重偏色和 halftone artifact。
+- 在目标 GPU 上 benchmark `attention_mechanism`。
+- 修改 `torch.compile` 后重新 smoke test，因为 compile cache 可能改变峰值显存。
+
+## 运行
+
+```bash
+simpletuner train example=minimaxh3-fl2va-convrot-int8.peft-lora
+```
+
+长训练前先跑短 smoke test，确认 `h3_drift_loss`、正常 loss 和验证样本走势一致。
diff --git a/documentation/quickstart/index.es.md b/documentation/quickstart/index.es.md
index b850bbbed..a281c78b8 100644
--- a/documentation/quickstart/index.es.md
+++ b/documentation/quickstart/index.es.md
@@ -61,6 +61,7 @@ Guías paso a paso para entrenar cada arquitectura de modelo compatible.
| **LTX Video 2** | 19B | [LTX-2 Community](https://ltx.io/model/license) | Aplican condiciones10 | [Guía de LTX Video 2](LTXVIDEO2.md) |
| **Cosmos3** | 4-65B | [OpenMDW 1.1](https://github.com/OpenMDW/openmdw/blob/main/1.1/LICENSE.OpenMDW-1.1) | Sí | [Guía de Cosmos3](COSMOS3.es.md) |
| **Hunyuan Video** | 8.3B | [Tencent Hunyuan Community](https://huggingface.co/tencent/HunyuanVideo-1.5/blob/main/LICENSE) | Aplican condiciones11 | [Guía de Hunyuan](HUNYUANVIDEO.md) |
+| **MiniMax H3** | 33B | [MiniMax H3 Community](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) | Aplican condiciones12 | [Guía de MiniMax H3](MINIMAX_H3.es.md) |
| **Sana Video** | - | [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) | Sí | [Guía de Sana Video](SANAVIDEO.md) |
| **Kandinsky 5 Video** | - | [MIT](https://opensource.org/license/mit) | Sí | [Guía de Kandinsky Video](KANDINSKY5_VIDEO.md) |
| **LongCat Video** | - | [MIT](https://opensource.org/license/mit) | Sí | [Guía de LongCat Video](LONGCAT_VIDEO.md) |
@@ -90,6 +91,8 @@ Guías paso a paso para entrenar cada arquitectura de modelo compatible.
11 La Tencent Hunyuan Community License incluye exclusiones territoriales y un umbral comercial para servicios muy grandes.
+12 La MiniMax H3 Community License excluye a Estados Unidos, la Unión Europea, Reino Unido y Corea del Sur del territorio aplicable estándar; esos territorios requieren autorización separada.
+
## Modelos de audio
diff --git a/documentation/quickstart/index.hi.md b/documentation/quickstart/index.hi.md
index 56a108939..6215b439b 100644
--- a/documentation/quickstart/index.hi.md
+++ b/documentation/quickstart/index.hi.md
@@ -61,6 +61,7 @@
| **LTX Video 2** | 19B | [LTX-2 Community](https://ltx.io/model/license) | शर्तें लागू10 | [LTX Video 2 गाइड](LTXVIDEO2.md) |
| **Cosmos3** | 4-65B | [OpenMDW 1.1](https://github.com/OpenMDW/openmdw/blob/main/1.1/LICENSE.OpenMDW-1.1) | हाँ | [Cosmos3 गाइड](COSMOS3.hi.md) |
| **Hunyuan Video** | 8.3B | [Tencent Hunyuan Community](https://huggingface.co/tencent/HunyuanVideo-1.5/blob/main/LICENSE) | शर्तें लागू11 | [Hunyuan गाइड](HUNYUANVIDEO.md) |
+| **MiniMax H3** | 33B | [MiniMax H3 Community](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) | शर्तें लागू12 | [MiniMax H3 गाइड](MINIMAX_H3.hi.md) |
| **Sana Video** | - | [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) | हाँ | [Sana Video गाइड](SANAVIDEO.md) |
| **Kandinsky 5 Video** | - | [MIT](https://opensource.org/license/mit) | हाँ | [Kandinsky Video गाइड](KANDINSKY5_VIDEO.md) |
| **LongCat Video** | - | [MIT](https://opensource.org/license/mit) | हाँ | [LongCat Video गाइड](LONGCAT_VIDEO.md) |
@@ -90,6 +91,8 @@
11 Tencent Hunyuan Community License में territorial exclusions और बहुत बड़े services के लिए commercial threshold शामिल है।
+12 MiniMax H3 Community License standard applicable territory से United States, European Union, United Kingdom, और South Korea को बाहर रखती है; उन territories में separate authorization चाहिए।
+
## ऑडियो मॉडल
diff --git a/documentation/quickstart/index.ja.md b/documentation/quickstart/index.ja.md
index 9c94fdfb0..c01a31a02 100644
--- a/documentation/quickstart/index.ja.md
+++ b/documentation/quickstart/index.ja.md
@@ -61,6 +61,7 @@
| **LTX Video 2** | 19B | [LTX-2 Community](https://ltx.io/model/license) | 条件付き10 | [LTX Video 2 ガイド](LTXVIDEO2.md) |
| **Cosmos3** | 4-65B | [OpenMDW 1.1](https://github.com/OpenMDW/openmdw/blob/main/1.1/LICENSE.OpenMDW-1.1) | はい | [Cosmos3 ガイド](COSMOS3.ja.md) |
| **Hunyuan Video** | 8.3B | [Tencent Hunyuan Community](https://huggingface.co/tencent/HunyuanVideo-1.5/blob/main/LICENSE) | 条件付き11 | [Hunyuan ガイド](HUNYUANVIDEO.md) |
+| **MiniMax H3** | 33B | [MiniMax H3 Community](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) | 条件付き12 | [MiniMax H3 ガイド](MINIMAX_H3.ja.md) |
| **Sana Video** | - | [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) | はい | [Sana Video ガイド](SANAVIDEO.md) |
| **Kandinsky 5 Video** | - | [MIT](https://opensource.org/license/mit) | はい | [Kandinsky Video ガイド](KANDINSKY5_VIDEO.md) |
| **LongCat Video** | - | [MIT](https://opensource.org/license/mit) | はい | [LongCat Video ガイド](LONGCAT_VIDEO.md) |
@@ -90,6 +91,8 @@
11 Tencent Hunyuan Community License には地域除外と、非常に大規模なサービス向けの商用しきい値があります。
+12 MiniMax H3 Community License は標準の適用地域から米国、欧州連合、英国、韓国を除外しており、これらの地域では別途認可が必要です。
+
## 音声モデル
diff --git a/documentation/quickstart/index.md b/documentation/quickstart/index.md
index df163cdc6..84d9d915d 100644
--- a/documentation/quickstart/index.md
+++ b/documentation/quickstart/index.md
@@ -61,6 +61,7 @@ Step-by-step guides for training each supported model architecture.
| **LTX Video 2** | 19B | [LTX-2 Community](https://ltx.io/model/license) | Conditions apply10 | [LTX Video 2 Guide](LTXVIDEO2.md) |
| **Cosmos3** | 4-65B | [OpenMDW 1.1](https://github.com/OpenMDW/openmdw/blob/main/1.1/LICENSE.OpenMDW-1.1) | Yes | [Cosmos3 Guide](COSMOS3.md) |
| **Hunyuan Video** | 8.3B | [Tencent Hunyuan Community](https://huggingface.co/tencent/HunyuanVideo-1.5/blob/main/LICENSE) | Conditions apply11 | [Hunyuan Guide](HUNYUANVIDEO.md) |
+| **MiniMax H3** | 33B | [MiniMax H3 Community](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) | Conditions apply12 | [MiniMax H3 Guide](MINIMAX_H3.md) |
| **Sana Video** | - | [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) | Yes | [Sana Video Guide](SANAVIDEO.md) |
| **Kandinsky 5 Video** | - | [MIT](https://opensource.org/license/mit) | Yes | [Kandinsky Video Guide](KANDINSKY5_VIDEO.md) |
| **LongCat Video** | - | [MIT](https://opensource.org/license/mit) | Yes | [LongCat Video Guide](LONGCAT_VIDEO.md) |
@@ -90,6 +91,8 @@ Step-by-step guides for training each supported model architecture.
11 Tencent Hunyuan Community License includes territorial exclusions and a commercial threshold for very large services.
+12 MiniMax H3 Community License excludes the United States, European Union, United Kingdom, and South Korea from the standard applicable territory; those territories require separate authorization.
+
## Audio Models
diff --git a/documentation/quickstart/index.pt-BR.md b/documentation/quickstart/index.pt-BR.md
index 1f37312d8..51b6fbd19 100644
--- a/documentation/quickstart/index.pt-BR.md
+++ b/documentation/quickstart/index.pt-BR.md
@@ -61,6 +61,7 @@ Guias passo a passo para treinar cada arquitetura de modelo suportada.
| **LTX Video 2** | 19B | [LTX-2 Community](https://ltx.io/model/license) | Condições aplicáveis10 | [Guia LTX Video 2](LTXVIDEO2.md) |
| **Cosmos3** | 4-65B | [OpenMDW 1.1](https://github.com/OpenMDW/openmdw/blob/main/1.1/LICENSE.OpenMDW-1.1) | Sim | [Guia Cosmos3](COSMOS3.pt-BR.md) |
| **Hunyuan Video** | 8.3B | [Tencent Hunyuan Community](https://huggingface.co/tencent/HunyuanVideo-1.5/blob/main/LICENSE) | Condições aplicáveis11 | [Guia Hunyuan](HUNYUANVIDEO.md) |
+| **MiniMax H3** | 33B | [MiniMax H3 Community](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) | Condicoes aplicaveis12 | [Guia MiniMax H3](MINIMAX_H3.pt-BR.md) |
| **Sana Video** | - | [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) | Sim | [Guia Sana Video](SANAVIDEO.md) |
| **Kandinsky 5 Video** | - | [MIT](https://opensource.org/license/mit) | Sim | [Guia Kandinsky Video](KANDINSKY5_VIDEO.md) |
| **LongCat Video** | - | [MIT](https://opensource.org/license/mit) | Sim | [Guia LongCat Video](LONGCAT_VIDEO.md) |
@@ -90,6 +91,8 @@ Guias passo a passo para treinar cada arquitetura de modelo suportada.
11 A Tencent Hunyuan Community License inclui exclusões territoriais e um limite comercial para serviços muito grandes.
+12 A MiniMax H3 Community License exclui Estados Unidos, União Europeia, Reino Unido e Coreia do Sul do território aplicável padrão; esses territórios exigem autorização separada.
+
## Modelos de Áudio
diff --git a/documentation/quickstart/index.zh.md b/documentation/quickstart/index.zh.md
index dc5271eb9..cb1a81fbe 100644
--- a/documentation/quickstart/index.zh.md
+++ b/documentation/quickstart/index.zh.md
@@ -61,6 +61,7 @@
| **LTX Video 2** | 19B | [LTX-2 Community](https://ltx.io/model/license) | 有条件10 | [LTX Video 2 指南](LTXVIDEO2.md) |
| **Cosmos3** | 4-65B | [OpenMDW 1.1](https://github.com/OpenMDW/openmdw/blob/main/1.1/LICENSE.OpenMDW-1.1) | 是 | [Cosmos3 指南](COSMOS3.zh.md) |
| **Hunyuan Video** | 8.3B | [Tencent Hunyuan Community](https://huggingface.co/tencent/HunyuanVideo-1.5/blob/main/LICENSE) | 有条件11 | [Hunyuan 指南](HUNYUANVIDEO.md) |
+| **MiniMax H3** | 33B | [MiniMax H3 Community](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) | 有条件12 | [MiniMax H3 指南](MINIMAX_H3.zh.md) |
| **Sana Video** | - | [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) | 是 | [Sana Video 指南](SANAVIDEO.md) |
| **Kandinsky 5 Video** | - | [MIT](https://opensource.org/license/mit) | 是 | [Kandinsky Video 指南](KANDINSKY5_VIDEO.md) |
| **LongCat Video** | - | [MIT](https://opensource.org/license/mit) | 是 | [LongCat Video 指南](LONGCAT_VIDEO.md) |
@@ -90,6 +91,8 @@
11 Tencent Hunyuan Community License 包含地域排除，以及针对超大规模服务的商用门槛。
+12 MiniMax H3 Community License 将美国、欧盟、英国和韩国排除在标准适用地区之外；这些地区需要单独授权。
+
## 音频模型
diff --git a/mkdocs.yml b/mkdocs.yml
index 8a4980ebb..f160ce0a5 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -276,6 +276,7 @@ nav:
- LTX Video 2: quickstart/LTXVIDEO2.md
- Cosmos3: quickstart/COSMOS3.md
- Hunyuan Video: quickstart/HUNYUANVIDEO.md
+ - MiniMax H3: quickstart/MINIMAX_H3.md
- Sana Video: quickstart/SANAVIDEO.md
- Kandinsky 5 Video: quickstart/KANDINSKY5_VIDEO.md
- LongCat Video: quickstart/LONGCAT_VIDEO.md
@@ -301,6 +302,7 @@ nav:
- TwinFlow (RCGM): distillation/TWINFLOW.md
- SDXL LCM: distillation/SDXL_LCM.md
- Wan DCM: distillation/WAN_DCM.md
+ - MiniMax H3 Drift: distillation/MINIMAX_H3_DRIFT.md
- FastWan DMD: distillation/FASTWAN_DMD.md
- GLANCE: distillation/GLANCE.md
- Evaluation:
diff --git a/simpletuner/diagnostics/__init__.py b/simpletuner/diagnostics/__init__.py
new file mode 100644
index 000000000..13b1f225d
--- /dev/null
+++ b/simpletuner/diagnostics/__init__.py
@@ -0,0 +1 @@
+"""Standalone diagnostics for inspecting SimpleTuner model behavior."""
diff --git a/simpletuner/diagnostics/h3_objective_geometry.py b/simpletuner/diagnostics/h3_objective_geometry.py
new file mode 100644
index 000000000..4ef46a9a9
--- /dev/null
+++ b/simpletuner/diagnostics/h3_objective_geometry.py
@@ -0,0 +1,565 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import random
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable
+
+import numpy as np
+import torch
+
+FLOWMAP_KEYS = (
+ "flowmap_r_timesteps",
+ "anyflow_r_timesteps",
+ "anyflow_timestep_interval",
+)
+
+
+@dataclass(frozen=True)
+class GeometryPoint:
+ timestep: float
+ sigma: float
+ model_timestep: float
+ r_timestep: float
+ r_sigma: float
+ anyflow_weight: float
+ drift_weight: float
+ sft_weight: float
+ normal_target: torch.Tensor
+ anyflow_target: torch.Tensor
+ flowmap_objective_target: torch.Tensor
+ base_prediction: torch.Tensor
+ drift_reference_prediction: torch.Tensor
+ normal_batch: dict[str, Any]
+ prepared_batch: dict[str, Any]
+
+
+def _flat_float(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.detach().float().reshape(-1)
+
+
+def tensor_norm(tensor: torch.Tensor) -> float:
+ return float(torch.linalg.vector_norm(_flat_float(tensor)).cpu())
+
+
+def cosine_similarity(left: torch.Tensor, right: torch.Tensor) -> float:
+ left_flat = _flat_float(left)
+ right_flat = _flat_float(right)
+ denominator = torch.linalg.vector_norm(left_flat) * torch.linalg.vector_norm(right_flat)
+ if float(denominator) == 0.0:
+ return float("nan")
+ return float(torch.dot(left_flat, right_flat).div(denominator).cpu())
+
+
+def norm_ratio(numerator: torch.Tensor, denominator: torch.Tensor) -> float:
+ denominator_norm = tensor_norm(denominator)
+ if denominator_norm == 0.0:
+ return float("nan")
+ return tensor_norm(numerator) / denominator_norm
+
+
+def mean_squared_error(left: torch.Tensor, right: torch.Tensor) -> float:
+ return float(torch.mean((_flat_float(left) - _flat_float(right)).square()).cpu())
+
+
+def trajectory_metrics(
+ *,
+ adapter_label: str,
+ point: GeometryPoint,
+ adapter_prediction: torch.Tensor,
+ normal_adapter_prediction: torch.Tensor,
+) -> dict[str, Any]:
+ adapter_residual = adapter_prediction - point.drift_reference_prediction
+ anyflow_correction = point.anyflow_target - point.drift_reference_prediction
+ objective_correction = point.flowmap_objective_target - point.drift_reference_prediction
+ normal_adapter_residual = normal_adapter_prediction - point.base_prediction
+ target_delta = point.anyflow_target - point.normal_target
+ flowmap_base_residual = point.drift_reference_prediction - point.base_prediction
+ return {
+ "adapter": adapter_label,
+ "timestep": point.timestep,
+ "sigma": point.sigma,
+ "model_timestep": point.model_timestep,
+ "r_timestep": point.r_timestep,
+ "r_sigma": point.r_sigma,
+ "interval": point.sigma - point.r_sigma,
+ "anyflow_weight": point.anyflow_weight,
+ "drift_weight": point.drift_weight,
+ "sft_weight": point.sft_weight,
+ "cos_adapter_base": cosine_similarity(adapter_prediction, point.base_prediction),
+ "adapter_base_norm_ratio": norm_ratio(adapter_prediction, point.base_prediction),
+ "cos_adapter_drift_reference": cosine_similarity(adapter_prediction, point.drift_reference_prediction),
+ "adapter_drift_reference_norm_ratio": norm_ratio(adapter_prediction, point.drift_reference_prediction),
+ "adapter_residual_norm": tensor_norm(adapter_residual),
+ "adapter_residual_base_norm_ratio": norm_ratio(adapter_residual, point.base_prediction),
+ "cos_adapter_residual_anyflow_correction": cosine_similarity(adapter_residual, anyflow_correction),
+ "adapter_residual_anyflow_correction_norm_ratio": norm_ratio(adapter_residual, anyflow_correction),
+ "cos_adapter_flowmap_objective": cosine_similarity(adapter_prediction, point.flowmap_objective_target),
+ "adapter_flowmap_objective_norm_ratio": norm_ratio(adapter_prediction, point.flowmap_objective_target),
+ "adapter_flowmap_objective_mse": mean_squared_error(adapter_prediction, point.flowmap_objective_target),
+ "cos_adapter_residual_objective_correction": cosine_similarity(adapter_residual, objective_correction),
+ "adapter_residual_objective_correction_norm_ratio": norm_ratio(adapter_residual, objective_correction),
+ "normal_adapter_residual_norm": tensor_norm(normal_adapter_residual),
+ "normal_adapter_residual_base_norm_ratio": norm_ratio(normal_adapter_residual, point.base_prediction),
+ "cos_normal_adapter_base": cosine_similarity(normal_adapter_prediction, point.base_prediction),
+ "normal_adapter_base_norm_ratio": norm_ratio(normal_adapter_prediction, point.base_prediction),
+ "cos_normal_adapter_normal_target": cosine_similarity(normal_adapter_prediction, point.normal_target),
+ "normal_adapter_normal_target_norm_ratio": norm_ratio(normal_adapter_prediction, point.normal_target),
+ "normal_adapter_normal_target_mse": mean_squared_error(normal_adapter_prediction, point.normal_target),
+ "cos_anyflow_normal_target": cosine_similarity(point.anyflow_target, point.normal_target),
+ "anyflow_normal_target_norm_ratio": norm_ratio(point.anyflow_target, point.normal_target),
+ "anyflow_target_delta_norm": tensor_norm(target_delta),
+ "anyflow_target_delta_normal_norm_ratio": norm_ratio(target_delta, point.normal_target),
+ "cos_base_normal_target": cosine_similarity(point.base_prediction, point.normal_target),
+ "base_normal_target_norm_ratio": norm_ratio(point.base_prediction, point.normal_target),
+ "base_normal_target_mse": mean_squared_error(point.base_prediction, point.normal_target),
+ "cos_drift_reference_anyflow_target": cosine_similarity(point.drift_reference_prediction, point.anyflow_target),
+ "drift_reference_anyflow_target_norm_ratio": norm_ratio(point.drift_reference_prediction, point.anyflow_target),
+ "drift_reference_anyflow_target_mse": mean_squared_error(point.drift_reference_prediction, point.anyflow_target),
+ "cos_drift_reference_base": cosine_similarity(point.drift_reference_prediction, point.base_prediction),
+ "flowmap_base_residual_norm": tensor_norm(flowmap_base_residual),
+ "flowmap_base_residual_base_norm_ratio": norm_ratio(flowmap_base_residual, point.base_prediction),
+ "adapter_prediction_norm": tensor_norm(adapter_prediction),
+ "base_prediction_norm": tensor_norm(point.base_prediction),
+ "drift_reference_prediction_norm": tensor_norm(point.drift_reference_prediction),
+ "normal_target_norm": tensor_norm(point.normal_target),
+ "anyflow_target_norm": tensor_norm(point.anyflow_target),
+ "flowmap_objective_target_norm": tensor_norm(point.flowmap_objective_target),
+ "normal_adapter_prediction_norm": tensor_norm(normal_adapter_prediction),
+ }
+
+
+def sample_vector(tensor: torch.Tensor, max_elements: int) -> np.ndarray:
+ flattened = _flat_float(tensor).cpu().numpy()
+ if max_elements <= 0 or flattened.size <= max_elements:
+ return flattened
+ indices = np.linspace(0, flattened.size - 1, num=max_elements, dtype=np.int64)
+ return flattened[indices]
+
+
+def pca_coordinates(vectors: Iterable[np.ndarray]) -> np.ndarray:
+ matrix = np.stack(list(vectors)).astype(np.float64, copy=False)
+ centered = matrix - matrix.mean(axis=0, keepdims=True)
+ if matrix.shape[0] < 2 or not np.any(centered):
+ return np.zeros((matrix.shape[0], 2), dtype=np.float64)
+ left, singular_values, _ = np.linalg.svd(centered, full_matrices=False)
+ coordinates = left[:, :2] * singular_values[:2]
+ if coordinates.shape[1] == 1:
+ coordinates = np.pad(coordinates, ((0, 0), (0, 1)))
+ return coordinates
+
+
+def _safe_name(value: str) -> str:
+ return re.sub(r"[^a-zA-Z0-9_.-]+", "_", value).strip("_") or "value"
+
+
+def _prediction_tensor(output: dict[str, Any]) -> torch.Tensor:
+ prediction = output.get("model_prediction")
+ if not torch.is_tensor(prediction):
+ raise ValueError("MiniMax-H3 diagnostic expected a tensor model_prediction.")
+ hidden_states_buffer = output.get("hidden_states_buffer")
+ if isinstance(hidden_states_buffer, dict):
+ hidden_states_buffer.clear()
+ return prediction.detach()
+
+
+def _without_flowmap_conditioning(batch: dict[str, Any]) -> dict[str, Any]:
+ normal_batch = dict(batch)
+ for key in FLOWMAP_KEYS:
+ normal_batch.pop(key, None)
+ normal_batch.pop("target", None)
+ normal_batch.pop("flow_target", None)
+ return normal_batch
+
+
+def _adapter_prediction(model, distiller, batch: dict[str, Any], *, enabled: bool) -> torch.Tensor:
+ distiller.toggle_adapter(enable=enabled)
+ try:
+ with torch.no_grad():
+ return _prediction_tensor(model.model_predict(batch)).float().cpu()
+ finally:
+ distiller.toggle_adapter(enable=True)
+
+
+def _seed_everything(seed: int) -> None:
+ random.seed(seed)
+ np.random.seed(seed % (2**32))
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+
+def _scalar_batch_value(batch: dict[str, Any], key: str, fallback: float) -> float:
+ value = batch.get(key)
+ if not torch.is_tensor(value):
+ return fallback
+ return float(value.detach().float().reshape(-1)[0].cpu())
+
+
+def prepare_geometry_points(trainer, prepared_batch: dict[str, Any], timesteps: list[float], seed: int):
+ distiller = trainer.distiller
+ inner_distiller = getattr(distiller, "inner_distiller", None)
+ if inner_distiller is None or inner_distiller.__class__.__name__ != "AnyFlowDistiller":
+ raise ValueError("H3 objective geometry requires h3_drift with inner_distillation_method=anyflow.")
+
+ points = []
+ anyflow_weight = float(inner_distiller.config.get("loss_weight", 1.0))
+ drift_weight = float(distiller.config.get("loss_weight", 1.0))
+ sft_weight = float(distiller.config.get("sft_loss_weight", 1.0))
+ flowmap_weight = anyflow_weight + drift_weight
+ if flowmap_weight <= 0.0:
+ raise ValueError("H3 objective geometry requires a positive AnyFlow or H3_DRIFT loss weight.")
+ for index, timestep in enumerate(timesteps):
+ normal_batch = trainer._prepare_custom_timestep_batch(prepared_batch, [timestep])
+ normal_batch = _without_flowmap_conditioning(normal_batch)
+ normal_target = distiller._normal_video_target(normal_batch).float().cpu()
+
+ _seed_everything(seed + index)
+ anyflow_batch = inner_distiller.prepare_batch(dict(normal_batch), trainer.model, trainer.state)
+ anyflow_target = anyflow_batch["target"].detach().float().cpu()
+
+ base_prediction = _adapter_prediction(trainer.model, distiller, normal_batch, enabled=False)
+ drift_reference = _adapter_prediction(trainer.model, distiller, anyflow_batch, enabled=False)
+ flowmap_objective_target = (anyflow_weight * anyflow_target + drift_weight * drift_reference) / flowmap_weight
+ sigma = _scalar_batch_value(normal_batch, "sigmas", timestep / 1000.0)
+ model_timestep = _scalar_batch_value(normal_batch, "timesteps", 1.0 - sigma)
+ r_timestep = _scalar_batch_value(anyflow_batch, "anyflow_r_timesteps", float("nan"))
+ points.append(
+ GeometryPoint(
+ timestep=float(timestep),
+ sigma=sigma,
+ model_timestep=model_timestep,
+ r_timestep=r_timestep,
+ r_sigma=1.0 - r_timestep,
+ anyflow_weight=anyflow_weight,
+ drift_weight=drift_weight,
+ sft_weight=sft_weight,
+ normal_target=normal_target,
+ anyflow_target=anyflow_target,
+ flowmap_objective_target=flowmap_objective_target,
+ base_prediction=base_prediction,
+ drift_reference_prediction=drift_reference,
+ normal_batch=normal_batch,
+ prepared_batch=anyflow_batch,
+ )
+ )
+ return points
+
+
+def _load_adapter(trainer, checkpoint_dir: Path) -> None:
+ component = trainer.model.get_trained_component(unwrap_model=False)
+ trainer.model.load_lora_weights([component], str(checkpoint_dir))
+
+
+def _first_raw_batch():
+ from simpletuner.helpers.data_backend.runtime import random_dataloader_iterator
+ from simpletuner.helpers.training.state_tracker import StateTracker
+
+ backends = {
+ backend_id: backend["train_dataloader"]
+ for backend_id, backend in StateTracker.get_data_backends().items()
+ if "train_dataloader" in backend and not StateTracker.backend_status(backend_id)
+ }
+ raw_batch = random_dataloader_iterator(1, backends)
+ if raw_batch is False or not isinstance(raw_batch, dict):
+ raise RuntimeError("Unable to fetch a diagnostic batch from the configured training data backend.")
+ return raw_batch
+
+
+def initialize_trainer(config: dict[str, Any]):
+ from simpletuner.helpers.training.attention_backend import AttentionBackendController, AttentionPhase
+ from simpletuner.helpers.training.trainer import Trainer
+
+ trainer = Trainer(config=config, exit_on_error=True)
+ trainer.init_noise_schedule()
+ trainer.init_seed()
+ trainer.init_huggingface_hub()
+ trainer.init_preprocessing_models()
+ trainer.init_precision(preprocessing_models_only=True)
+ trainer.init_data_backend()
+ trainer.init_unload_text_encoder()
+ trainer.init_unload_vae()
+ trainer.init_load_base_model()
+ trainer.init_delete_model_caches()
+ trainer.init_controlnet_model()
+ trainer.init_tread_model()
+ trainer.init_precision()
+ trainer.init_freeze_models()
+ trainer.init_trainable_peft_adapter()
+ trainer.move_models(destination="accelerator")
+ trainer.init_distillation()
+ AttentionBackendController.apply(trainer.config, AttentionPhase.TRAIN)
+ trainer.model.get_trained_component(unwrap_model=False).eval()
+ return trainer
+
+
+def _diagnostic_config(config_path: Path, output_dir: Path) -> dict[str, Any]:
+ with config_path.open("r", encoding="utf-8") as handle:
+ config = json.load(handle)
+ config.update(
+ {
+ "checkpoint_step_interval": 0,
+ "dataloader_prefetch": False,
+ "max_train_steps": 1,
+ "output_dir": str(output_dir / "trainer-output"),
+ "push_checkpoints_to_hub": False,
+ "push_to_hub": False,
+ "report_to": "none",
+ "resume_from_checkpoint": None,
+ "validation_on_startup": False,
+ "validation_step_interval": 0,
+ "validation_steps": 0,
+ }
+ )
+ return config
+
+
+def _parse_adapter(value: str) -> tuple[str, Path]:
+ if "=" not in value:
+ raise argparse.ArgumentTypeError("Adapters must use LABEL=/path/to/checkpoint syntax.")
+ label, path = value.split("=", 1)
+ if not label.strip() or not path.strip():
+ raise argparse.ArgumentTypeError("Adapters must have a non-empty label and path.")
+ return label.strip(), Path(path).expanduser()
+
+
+def _parse_timesteps(value: str) -> list[float]:
+ try:
+ timesteps = [float(item.strip()) for item in value.split(",") if item.strip()]
+ except ValueError as exc:
+ raise argparse.ArgumentTypeError("Timesteps must be comma-separated numbers.") from exc
+ if not timesteps or any(timestep <= 0.0 or timestep > 1000.0 for timestep in timesteps):
+ raise argparse.ArgumentTypeError("Timesteps must be in (0, 1000].")
+ return timesteps
+
+
+def _write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
+ with path.open("w", newline="", encoding="utf-8") as handle:
+ writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
+ writer.writeheader()
+ writer.writerows(rows)
+
+
+def _write_npz(
+ path: Path,
+ points: list[GeometryPoint],
+ adapter_vectors: dict[tuple[str, float, str], torch.Tensor],
+ max_vector_elements: int,
+) -> tuple[list[str], list[np.ndarray]]:
+ arrays: dict[str, np.ndarray] = {}
+ pca_labels: list[str] = []
+ pca_vectors: list[np.ndarray] = []
+ for point in points:
+ timestep_name = f"t{point.timestep:g}"
+ for kind, tensor in (
+ ("normal_target", point.normal_target),
+ ("anyflow_target", point.anyflow_target),
+ ("flowmap_objective_target", point.flowmap_objective_target),
+ ("base_prediction", point.base_prediction),
+ ("drift_reference", point.drift_reference_prediction),
+ ):
+ label = f"{kind}:{timestep_name}"
+ vector = sample_vector(tensor, max_vector_elements)
+ arrays[_safe_name(label)] = vector
+ pca_labels.append(label)
+ pca_vectors.append(vector)
+ for (adapter_label, timestep, branch), tensor in adapter_vectors.items():
+ label = f"adapter_prediction:{branch}:{adapter_label}:t{timestep:g}"
+ vector = sample_vector(tensor, max_vector_elements)
+ arrays[_safe_name(label)] = vector
+ pca_labels.append(label)
+ pca_vectors.append(vector)
+ arrays["metadata_json"] = np.asarray(json.dumps({"pca_labels": pca_labels, "max_vector_elements": max_vector_elements}))
+ np.savez_compressed(path, **arrays)
+ return pca_labels, pca_vectors
+
+
+def _write_plots(
+ output_dir: Path,
+ rows: list[dict[str, Any]],
+ pca_labels: list[str],
+ pca_vectors: list[np.ndarray],
+) -> None:
+ try:
+ import matplotlib
+
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+ except ImportError as exc:
+ raise RuntimeError("matplotlib is required to write H3 geometry plots.") from exc
+
+ coordinates = pca_coordinates(pca_vectors)
+ figure, axis = plt.subplots(figsize=(11, 8))
+ groups: dict[str, list[int]] = {}
+ for index, label in enumerate(pca_labels):
+ group = label.split(":", 1)[0]
+ groups.setdefault(group, []).append(index)
+ for group, indices in groups.items():
+ axis.scatter(coordinates[indices, 0], coordinates[indices, 1], label=group, s=24)
+ axis.set_title("MiniMax-H3 objective trajectory PCA")
+ axis.set_xlabel("PC1")
+ axis.set_ylabel("PC2")
+ axis.legend(fontsize=8)
+ axis.grid(alpha=0.2)
+ figure.tight_layout()
+ figure.savefig(output_dir / "pca_trajectory.png", dpi=160)
+ plt.close(figure)
+
+ figure, axes = plt.subplots(2, 1, figsize=(11, 9), sharex=True)
+ adapter_labels = list(dict.fromkeys(str(row["adapter"]) for row in rows))
+ for adapter_label in adapter_labels:
+ adapter_rows = sorted(
+ (row for row in rows if row["adapter"] == adapter_label), key=lambda row: float(row["timestep"])
+ )
+ timesteps = [float(row["timestep"]) for row in adapter_rows]
+ axes[0].plot(
+ timesteps,
+ [float(row["cos_adapter_flowmap_objective"]) for row in adapter_rows],
+ marker="o",
+ label=f"{adapter_label}: FlowMap/objective cosine",
+ )
+ axes[1].plot(
+ timesteps,
+ [float(row["adapter_residual_base_norm_ratio"]) for row in adapter_rows],
+ marker="o",
+ label=f"{adapter_label}: residual/base norm",
+ )
+ axes[1].plot(
+ timesteps,
+ [float(row["normal_adapter_residual_base_norm_ratio"]) for row in adapter_rows],
+ marker="x",
+ linestyle=":",
+ label=f"{adapter_label}: normal residual/base norm",
+ )
+ first_by_timestep = {}
+ for row in rows:
+ first_by_timestep.setdefault(float(row["timestep"]), row)
+ target_rows = [first_by_timestep[key] for key in sorted(first_by_timestep)]
+ target_timesteps = [float(row["timestep"]) for row in target_rows]
+ axes[0].plot(
+ target_timesteps,
+ [float(row["cos_anyflow_normal_target"]) for row in target_rows],
+ color="black",
+ linestyle="--",
+ label="AnyFlow/normal target cosine",
+ )
+ axes[1].plot(
+ target_timesteps,
+ [float(row["anyflow_normal_target_norm_ratio"]) for row in target_rows],
+ color="black",
+ linestyle="--",
+ label="AnyFlow/normal target norm",
+ )
+ axes[0].set_ylabel("Cosine similarity")
+ axes[1].set_ylabel("Norm ratio")
+ axes[1].set_xlabel("Timestep")
+ for axis in axes:
+ axis.grid(alpha=0.2)
+ axis.legend(fontsize=8)
+ figure.suptitle("MiniMax-H3 norm and cosine geometry")
+ figure.tight_layout()
+ figure.savefig(output_dir / "norm_cosine_by_timestep.png", dpi=160)
+ plt.close(figure)
+
+
+def run_diagnostic(args: argparse.Namespace) -> None:
+ args.output_dir.mkdir(parents=True, exist_ok=True)
+ config = _diagnostic_config(args.config, args.output_dir)
+ trainer = initialize_trainer(config)
+ try:
+ _seed_everything(args.seed)
+ raw_batch = _first_raw_batch()
+ prepared_batch = trainer.model.prepare_batch(raw_batch, state=trainer.state)
+ points = prepare_geometry_points(trainer, prepared_batch, args.timesteps, args.seed)
+
+ rows = []
+ adapter_vectors: dict[tuple[str, float, str], torch.Tensor] = {}
+ adapters = [("base", None), *args.adapter]
+ for adapter_label, checkpoint_dir in adapters:
+ if checkpoint_dir is not None:
+ _load_adapter(trainer, checkpoint_dir)
+ for point in points:
+ adapter_prediction = _adapter_prediction(
+ trainer.model,
+ trainer.distiller,
+ point.prepared_batch,
+ enabled=checkpoint_dir is not None,
+ )
+ normal_adapter_prediction = _adapter_prediction(
+ trainer.model,
+ trainer.distiller,
+ point.normal_batch,
+ enabled=checkpoint_dir is not None,
+ )
+ rows.append(
+ trajectory_metrics(
+ adapter_label=adapter_label,
+ point=point,
+ adapter_prediction=adapter_prediction,
+ normal_adapter_prediction=normal_adapter_prediction,
+ )
+ )
+ adapter_vectors[(adapter_label, point.timestep, "flowmap")] = adapter_prediction
+ adapter_vectors[(adapter_label, point.timestep, "normal")] = normal_adapter_prediction
+
+ _write_csv(args.output_dir / "trajectory_metrics.csv", rows)
+ pca_labels, pca_vectors = _write_npz(
+ args.output_dir / "trajectory_vectors.npz",
+ points,
+ adapter_vectors,
+ args.max_vector_elements,
+ )
+ _write_plots(args.output_dir, rows, pca_labels, pca_vectors)
+ with (args.output_dir / "run_metadata.json").open("w", encoding="utf-8") as handle:
+ json.dump(
+ {
+ "config": str(args.config),
+ "adapters": [(label, None if path is None else str(path)) for label, path in adapters],
+ "timesteps": args.timesteps,
+ "seed": args.seed,
+ "max_vector_elements": args.max_vector_elements,
+ "batch": {
+ "data_backend_id": prepared_batch.get("data_backend_id"),
+ "filepaths": [str(path) for path in prepared_batch.get("filepaths", [])],
+ },
+ },
+ handle,
+ indent=2,
+ )
+ finally:
+ trainer.cleanup()
+
+
+def build_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(description="Inspect MiniMax-H3 AnyFlow/H3_DRIFT objective geometry.")
+ parser.add_argument("--config", type=Path, required=True, help="SimpleTuner config.json path.")
+ parser.add_argument(
+ "--adapter",
+ action="append",
+ default=[],
+ type=_parse_adapter,
+ metavar="LABEL=CHECKPOINT_DIR",
+ help="Checkpoint adapter to compare; may be repeated.",
+ )
+ parser.add_argument("--output-dir", type=Path, required=True)
+ parser.add_argument(
+ "--timesteps",
+ type=_parse_timesteps,
+ default=_parse_timesteps("50,100,250,500,750,900,975"),
+ )
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--max-vector-elements", type=int, default=65536)
+ return parser
+
+
+def main() -> None:
+ args = build_parser().parse_args()
+ run_diagnostic(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch/config.json b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch/config.json
new file mode 100644
index 000000000..7056d687b
--- /dev/null
+++ b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch/config.json
@@ -0,0 +1,78 @@
+{
+ "attention_mechanism": "native-efficient",
+ "audio_flow_schedule_shift": 3.0,
+ "base_model_precision": "no_change",
+ "caption_dropout_probability": 0.1,
+ "checkpoint_step_interval": 100,
+ "checkpoints_total_limit": 5,
+ "compress_disk_cache": true,
+ "data_backend_config": "simpletuner/examples/multidatabackend-minimaxh3-disney-fl2va-480p+39f.json",
+ "disable_benchmark": false,
+ "disable_bucket_pruning": true,
+ "distillation_config": {
+ "h3_drift": {
+ "audio_weight": 1.0,
+ "balance": "token",
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "video_weight": 1.0
+ }
+ },
+ "distillation_method": "h3_drift",
+ "flow_schedule_shift": 12.0,
+ "gradient_accumulation_steps": 1,
+ "gradient_checkpointing": true,
+ "gradient_checkpointing_backend": "torch",
+ "hub_model_id": "simpletuner-minimaxh3-fl2va-convrot-int8-24g-ramtorch-lora-test",
+ "learning_rate": 6e-5,
+ "lora_rank": 16,
+ "lora_type": "standard",
+ "lr_scheduler": "constant_with_warmup",
+ "lr_warmup_steps": 50,
+ "max_grad_norm": 0.1,
+ "max_train_steps": 200,
+ "minimum_image_size": 0,
+ "mixed_precision": "bf16",
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "musubi_blocks_to_swap": 0,
+ "num_processes": 1,
+ "num_train_epochs": 0,
+ "offload_during_startup": true,
+ "optimizer": "adamw_bf16",
+ "output_dir": "output/examples/minimaxh3-fl2va-convrot-int8-24g.peft-lora+ramtorch",
+ "push_checkpoints_to_hub": false,
+ "push_to_hub": false,
+ "quantize_via": "cpu",
+ "ramtorch": true,
+ "ramtorch_disable_extensions": true,
+ "ramtorch_disable_sync_hooks": false,
+ "ramtorch_target_modules": "*",
+ "ramtorch_text_encoder": false,
+ "ramtorch_transformer_percent": 100,
+ "report_to": "none",
+ "resolution": 480,
+ "resolution_type": "pixel_area",
+ "seed": 42,
+ "text_encoder_1_precision": "int8-quanto",
+ "tracker_project_name": "lora-training",
+ "tracker_run_name": "example-training-run-minimaxh3-fl2va-convrot-int8-24g-ramtorch",
+ "train_batch_size": 1,
+ "use_ema": false,
+ "vae_batch_size": 1,
+ "vae_enable_slicing": true,
+ "vae_enable_temporal_roll": true,
+ "vae_enable_tiling": true,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0,
+ "validation_negative_prompt": "",
+ "validation_num_inference_steps": 30,
+ "validation_num_video_frames": 39,
+ "validation_prompt": "A cinematic video of a small robot holding a sign that says hello world from MiniMax H3.",
+ "validation_prompt_library": false,
+ "validation_resolution": "480x480",
+ "validation_seed": 42,
+ "validation_steps": 50
+}
diff --git a/simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora/config.json b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora/config.json
new file mode 100644
index 000000000..2c50e8994
--- /dev/null
+++ b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora/config.json
@@ -0,0 +1,73 @@
+{
+ "attention_mechanism": "native-efficient",
+ "audio_flow_schedule_shift": 3.0,
+ "base_model_precision": "no_change",
+ "caption_dropout_probability": 0.1,
+ "checkpoint_step_interval": 100,
+ "checkpoints_total_limit": 5,
+ "compress_disk_cache": true,
+ "data_backend_config": "simpletuner/examples/multidatabackend-minimaxh3-disney-fl2va-480p+39f.json",
+ "disable_benchmark": false,
+ "disable_bucket_pruning": true,
+ "distillation_config": {
+ "h3_drift": {
+ "audio_weight": 1.0,
+ "balance": "token",
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "video_weight": 1.0
+ }
+ },
+ "distillation_method": "h3_drift",
+ "flow_schedule_shift": 12.0,
+ "gradient_accumulation_steps": 1,
+ "gradient_checkpointing": true,
+ "gradient_checkpointing_backend": "torch",
+ "gradient_checkpointing_interval": 2,
+ "hub_model_id": "simpletuner-minimaxh3-fl2va-convrot-int8-32g-lora-test",
+ "learning_rate": 6e-5,
+ "lora_rank": 16,
+ "lora_type": "standard",
+ "lr_scheduler": "constant_with_warmup",
+ "lr_warmup_steps": 50,
+ "max_grad_norm": 0.1,
+ "max_train_steps": 200,
+ "minimum_image_size": 0,
+ "mixed_precision": "bf16",
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "musubi_blocks_to_swap": 0,
+ "num_processes": 1,
+ "num_train_epochs": 0,
+ "offload_during_startup": true,
+ "optimizer": "adamw_bf16",
+ "output_dir": "output/examples/minimaxh3-fl2va-convrot-int8-32g.peft-lora",
+ "push_checkpoints_to_hub": false,
+ "push_to_hub": false,
+ "quantize_via": "cpu",
+ "report_to": "none",
+ "resolution": 480,
+ "resolution_type": "pixel_area",
+ "seed": 42,
+ "text_encoder_1_precision": "int8-quanto",
+ "tracker_project_name": "lora-training",
+ "tracker_run_name": "example-training-run-minimaxh3-fl2va-convrot-int8-32g",
+ "train_batch_size": 1,
+ "use_ema": false,
+ "vae_batch_size": 1,
+ "vae_enable_slicing": true,
+ "vae_enable_temporal_roll": true,
+ "vae_enable_tiling": true,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0,
+ "validation_negative_prompt": "",
+ "validation_num_inference_steps": 30,
+ "validation_num_video_frames": 39,
+ "validation_prompt": "A cinematic video of a small robot holding a sign that says hello world from MiniMax H3.",
+ "validation_prompt_library": false,
+ "validation_resolution": "480x480",
+ "validation_seed": 42,
+ "validation_steps": 50
+}
diff --git a/simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora/config.json b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora/config.json
new file mode 100644
index 000000000..8c8a47c33
--- /dev/null
+++ b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora/config.json
@@ -0,0 +1,73 @@
+{
+ "attention_mechanism": "native-efficient",
+ "audio_flow_schedule_shift": 3.0,
+ "base_model_precision": "no_change",
+ "caption_dropout_probability": 0.1,
+ "checkpoint_step_interval": 100,
+ "checkpoints_total_limit": 5,
+ "compress_disk_cache": true,
+ "data_backend_config": "simpletuner/examples/multidatabackend-minimaxh3-disney-fl2va-480p+39f.json",
+ "disable_benchmark": false,
+ "disable_bucket_pruning": true,
+ "distillation_config": {
+ "h3_drift": {
+ "audio_weight": 1.0,
+ "balance": "token",
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "video_weight": 1.0
+ }
+ },
+ "distillation_method": "h3_drift",
+ "flow_schedule_shift": 12.0,
+ "gradient_accumulation_steps": 1,
+ "gradient_checkpointing": true,
+ "gradient_checkpointing_backend": "torch",
+ "gradient_checkpointing_interval": 2,
+ "hub_model_id": "simpletuner-minimaxh3-fl2va-convrot-int8-48g-lora-test",
+ "learning_rate": 6e-5,
+ "lora_rank": 16,
+ "lora_type": "standard",
+ "lr_scheduler": "constant_with_warmup",
+ "lr_warmup_steps": 50,
+ "max_grad_norm": 0.1,
+ "max_train_steps": 200,
+ "minimum_image_size": 0,
+ "mixed_precision": "bf16",
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "musubi_blocks_to_swap": 0,
+ "num_processes": 1,
+ "num_train_epochs": 0,
+ "offload_during_startup": true,
+ "optimizer": "adamw_bf16",
+ "output_dir": "output/examples/minimaxh3-fl2va-convrot-int8-48g.peft-lora",
+ "push_checkpoints_to_hub": false,
+ "push_to_hub": false,
+ "quantize_via": "cpu",
+ "report_to": "none",
+ "resolution": 480,
+ "resolution_type": "pixel_area",
+ "seed": 42,
+ "text_encoder_1_precision": "int8-quanto",
+ "tracker_project_name": "lora-training",
+ "tracker_run_name": "example-training-run-minimaxh3-fl2va-convrot-int8-48g",
+ "train_batch_size": 1,
+ "use_ema": false,
+ "vae_batch_size": 1,
+ "vae_enable_slicing": true,
+ "vae_enable_temporal_roll": true,
+ "vae_enable_tiling": true,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0,
+ "validation_negative_prompt": "",
+ "validation_num_inference_steps": 30,
+ "validation_num_video_frames": 39,
+ "validation_prompt": "A cinematic video of a small robot holding a sign that says hello world from MiniMax H3.",
+ "validation_prompt_library": false,
+ "validation_resolution": "480x480",
+ "validation_seed": 42,
+ "validation_steps": 50
+}
diff --git a/simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora/config.json b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora/config.json
new file mode 100644
index 000000000..cf00a480a
--- /dev/null
+++ b/simpletuner/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora/config.json
@@ -0,0 +1,71 @@
+{
+ "attention_mechanism": "native-efficient",
+ "audio_flow_schedule_shift": 3.0,
+ "base_model_precision": "no_change",
+ "caption_dropout_probability": 0.1,
+ "checkpoint_step_interval": 100,
+ "checkpoints_total_limit": 5,
+ "compress_disk_cache": true,
+ "data_backend_config": "simpletuner/examples/multidatabackend-minimaxh3-disney-fl2va-480p+39f.json",
+ "disable_benchmark": false,
+ "disable_bucket_pruning": true,
+ "distillation_config": {
+ "h3_drift": {
+ "audio_weight": 1.0,
+ "balance": "token",
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "video_weight": 1.0
+ }
+ },
+ "distillation_method": "h3_drift",
+ "flow_schedule_shift": 12.0,
+ "gradient_accumulation_steps": 1,
+ "gradient_checkpointing": false,
+ "hub_model_id": "simpletuner-minimaxh3-fl2va-convrot-int8-80g-lora-test",
+ "learning_rate": 6e-5,
+ "lora_rank": 16,
+ "lora_type": "standard",
+ "lr_scheduler": "constant_with_warmup",
+ "lr_warmup_steps": 50,
+ "max_grad_norm": 0.1,
+ "max_train_steps": 200,
+ "minimum_image_size": 0,
+ "mixed_precision": "bf16",
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "musubi_blocks_to_swap": 0,
+ "num_processes": 1,
+ "num_train_epochs": 0,
+ "offload_during_startup": true,
+ "optimizer": "adamw_bf16",
+ "output_dir": "output/examples/minimaxh3-fl2va-convrot-int8-80g.peft-lora",
+ "push_checkpoints_to_hub": false,
+ "push_to_hub": false,
+ "quantize_via": "cpu",
+ "report_to": "none",
+ "resolution": 480,
+ "resolution_type": "pixel_area",
+ "seed": 42,
+ "text_encoder_1_precision": "int8-quanto",
+ "tracker_project_name": "lora-training",
+ "tracker_run_name": "example-training-run-minimaxh3-fl2va-convrot-int8-80g",
+ "train_batch_size": 1,
+ "use_ema": false,
+ "vae_batch_size": 1,
+ "vae_enable_slicing": true,
+ "vae_enable_temporal_roll": true,
+ "vae_enable_tiling": true,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0,
+ "validation_negative_prompt": "",
+ "validation_num_inference_steps": 30,
+ "validation_num_video_frames": 39,
+ "validation_prompt": "A cinematic video of a small robot holding a sign that says hello world from MiniMax H3.",
+ "validation_prompt_library": false,
+ "validation_resolution": "480x480",
+ "validation_seed": 42,
+ "validation_steps": 50
+}
diff --git a/simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora/config.json b/simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora/config.json
new file mode 100644
index 000000000..569571c5e
--- /dev/null
+++ b/simpletuner/examples/minimaxh3-fl2va-convrot-int8.peft-lora/config.json
@@ -0,0 +1,72 @@
+{
+ "attention_mechanism": "native-efficient",
+ "audio_flow_schedule_shift": 3.0,
+ "base_model_precision": "no_change",
+ "caption_dropout_probability": 0.1,
+ "checkpoint_step_interval": 100,
+ "checkpoints_total_limit": 5,
+ "compress_disk_cache": true,
+ "data_backend_config": "simpletuner/examples/multidatabackend-minimaxh3-disney-fl2va-480p+39f.json",
+ "disable_benchmark": false,
+ "disable_bucket_pruning": true,
+ "distillation_config": {
+ "h3_drift": {
+ "audio_weight": 1.0,
+ "balance": "token",
+ "loss_weight": 0.5,
+ "sft_loss_weight": 1.0,
+ "video_weight": 1.0
+ }
+ },
+ "distillation_method": "h3_drift",
+ "flow_schedule_shift": 12.0,
+ "gradient_accumulation_steps": 1,
+ "gradient_checkpointing": true,
+ "gradient_checkpointing_backend": "torch",
+ "hub_model_id": "simpletuner-minimaxh3-fl2va-convrot-int8-lora-test",
+ "learning_rate": 6e-5,
+ "lora_rank": 16,
+ "lora_type": "standard",
+ "lr_scheduler": "constant_with_warmup",
+ "lr_warmup_steps": 50,
+ "max_grad_norm": 0.1,
+ "max_train_steps": 200,
+ "minimum_image_size": 0,
+ "mixed_precision": "bf16",
+ "model_family": "minimaxh3",
+ "model_flavour": "convrot-int8",
+ "model_type": "lora",
+ "musubi_blocks_to_swap": 25,
+ "num_processes": 1,
+ "num_train_epochs": 0,
+ "offload_during_startup": true,
+ "optimizer": "adamw_bf16",
+ "output_dir": "output/examples/minimaxh3-fl2va-convrot-int8.peft-lora",
+ "push_checkpoints_to_hub": false,
+ "push_to_hub": false,
+ "quantize_via": "cpu",
+ "report_to": "none",
+ "resolution": 480,
+ "resolution_type": "pixel_area",
+ "seed": 42,
+ "text_encoder_1_precision": "int8-quanto",
+ "tracker_project_name": "lora-training",
+ "tracker_run_name": "example-training-run-minimaxh3-fl2va-convrot-int8",
+ "train_batch_size": 1,
+ "use_ema": false,
+ "vae_batch_size": 1,
+ "vae_enable_slicing": true,
+ "vae_enable_temporal_roll": true,
+ "vae_enable_tiling": true,
+ "validation_disable_unconditional": true,
+ "validation_guidance": 1.0,
+ "validation_guidance_real": 1.0,
+ "validation_negative_prompt": "",
+ "validation_num_inference_steps": 30,
+ "validation_num_video_frames": 39,
+ "validation_prompt": "A cinematic video of a small robot holding a sign that says hello world from MiniMax H3.",
+ "validation_prompt_library": false,
+ "validation_resolution": "480x480",
+ "validation_seed": 42,
+ "validation_steps": 50
+}
diff --git a/simpletuner/helpers/caching/text_embeds.py b/simpletuner/helpers/caching/text_embeds.py
index f4923dd57..103c26b3d 100644
--- a/simpletuner/helpers/caching/text_embeds.py
+++ b/simpletuner/helpers/caching/text_embeds.py
@@ -19,6 +19,7 @@
from simpletuner.helpers.training.multi_process import rank_info
from simpletuner.helpers.training.state_tracker import StateTracker
from simpletuner.helpers.training.wrappers import gather_dict_of_tensors_shapes, move_dict_of_tensors_to_device
+from simpletuner.helpers.utils.pathing import canonicalize_data_uri
from simpletuner.helpers.webhooks.mixin import WebhookMixin
logger = logging.getLogger("TextEmbeddingCache")
@@ -129,8 +130,9 @@ def _normalize_key_value(self, key_value):
if "://" not in normalized:
normalized = os.path.normcase(os.path.abspath(os.path.normpath(normalized)))
elif self.key_type is TextEmbedCacheKey.DATASET_AND_FILENAME:
- # Keys already include dataset identifiers; leave as-is.
- pass
+ dataset_id, separator, data_path = normalized.partition(":")
+ if separator:
+ normalized = f"{dataset_id}:{canonicalize_data_uri(data_path)}"
return normalized
def create_hash(self, key_value):
@@ -258,11 +260,15 @@ def _encode_and_cache_prompt_batch(
pending_records: List[PromptCacheRecord],
pending_filenames: List[str],
is_negative_prompt: bool = False,
+ is_validation: bool = False,
) -> None:
prompts = [record.get("prompt") for record in pending_records]
prompt_contexts = [record.get("metadata") or {} for record in pending_records]
- text_encoder_output = self.model.encode_text_batch(
- prompts, is_negative_prompt=is_negative_prompt, prompt_contexts=prompt_contexts
+ text_encoder_output = self._encode_text_batch(
+ prompts,
+ is_negative_prompt=is_negative_prompt,
+ prompt_contexts=prompt_contexts,
+ is_validation=is_validation,
)
if not isinstance(text_encoder_output, dict):
raise TypeError(f"encode_text_batch returned {type(text_encoder_output)}; expected a dict.")
@@ -286,6 +292,32 @@ def _encode_and_cache_prompt_batch(
self._slice_batch_output_for_cache(text_encoder_output, index, batch_size),
)
+ def _encode_text_batch(
+ self,
+ prompts: list,
+ *,
+ is_negative_prompt: bool,
+ prompt_contexts: Optional[List[dict]],
+ is_validation: bool,
+ ):
+ had_validation_marker = hasattr(self.model, "_current_prompt_is_validation")
+ previous_validation_marker = getattr(self.model, "_current_prompt_is_validation", None)
+ self.model._current_prompt_is_validation = bool(is_validation)
+ try:
+ return self.model.encode_text_batch(
+ prompts,
+ is_negative_prompt=is_negative_prompt,
+ prompt_contexts=prompt_contexts,
+ )
+ finally:
+ if had_validation_marker:
+ self.model._current_prompt_is_validation = previous_validation_marker
+ else:
+ try:
+ delattr(self.model, "_current_prompt_is_validation")
+ except AttributeError:
+ pass
+
def discover_all_files(self):
"""Identify all files in the data backend."""
logger.info(f"{self.rank_info}(id={self.id}) Listing all text embed cache entries")
@@ -455,6 +487,7 @@ def compute_embeddings_for_prompts(
load_from_cache: bool = True,
is_negative_prompt: bool = False,
progress_callback=None,
+ split_between_processes: bool = True,
):
if self.model is None:
self.model = StateTracker.get_model()
@@ -512,7 +545,7 @@ def compute_embeddings_for_prompts(
should_encode = not load_from_cache
if requires_context and uncached_records and should_encode:
contextless_records = [record for record in raw_records if not record.get("metadata")]
- if contextless_records:
+ if contextless_records and not is_validation:
logger.warning(
f"{self.rank_info}(id={self.id}) Prompt records require image context but metadata "
f"was missing for {len(contextless_records)} entries. "
@@ -523,6 +556,10 @@ def compute_embeddings_for_prompts(
raw_records = [record for record in raw_records if record.get("metadata")]
if not raw_records:
return None
+ elif contextless_records:
+ self.debug_log(
+ "Encoding contextless validation prompt(s) for a model that normally requires image metadata."
+ )
output = None
if self.model is not None:
@@ -533,6 +570,7 @@ def compute_embeddings_for_prompts(
load_from_cache=load_from_cache,
is_negative_prompt=is_negative_prompt,
progress_callback=progress_callback,
+ split_between_processes=split_between_processes,
)
else:
raise ValueError(f"No such text encoding backend for model type '{self.model_type}'")
@@ -557,12 +595,13 @@ def compute_prompt_embeddings_with_model(
load_from_cache: bool = True,
is_negative_prompt: bool = False,
progress_callback=None,
+ split_between_processes: bool = True,
):
prompt_embeds_all = []
should_encode = not load_from_cache
args = StateTracker.get_args()
records = self._normalize_prompt_records(prompt_records) if prompt_records is not None else self.prompt_records
- if should_encode:
+ if should_encode and split_between_processes:
local_records = self.split_prompt_records_between_processes(records)
else:
local_records = records
@@ -652,6 +691,7 @@ def compute_prompt_embeddings_with_model(
pending_records,
pending_filenames,
is_negative_prompt=is_negative_prompt,
+ is_validation=is_validation,
)
pending_records = []
pending_filenames = []
@@ -661,8 +701,11 @@ def compute_prompt_embeddings_with_model(
encoder_device = self.text_encoders[0].device if self.text_encoders else "model"
self.debug_log(f"Encoding filename {filename} :: device {encoder_device} :: prompt {prompt}")
prompt_contexts = [record.get("metadata") or {}]
- text_encoder_output = self.model.encode_text_batch(
- [prompt], is_negative_prompt=is_negative_prompt, prompt_contexts=prompt_contexts
+ text_encoder_output = self._encode_text_batch(
+ [prompt],
+ is_negative_prompt=is_negative_prompt,
+ prompt_contexts=prompt_contexts,
+ is_validation=is_validation,
)
logger.debug(
f"Filename {filename} prompt embeds: {gather_dict_of_tensors_shapes(tensors=text_encoder_output)}, keys: {text_encoder_output.keys()}"
@@ -710,6 +753,7 @@ def compute_prompt_embeddings_with_model(
pending_records,
pending_filenames,
is_negative_prompt=is_negative_prompt,
+ is_validation=is_validation,
)
while self.write_queue.qsize() > 0:
@@ -736,6 +780,10 @@ def compute_prompt_embeddings_with_model(
if len(prompt_embeds_all) == 1:
return prompt_embeds_all[0]
+ transformed_encoder_output = self.model.collate_prompt_embeds(prompt_embeds_all)
+ if transformed_encoder_output:
+ return transformed_encoder_output
+
# Concatenate multiple embeddings
if prompt_embeds_all and "prompt_embeds" in prompt_embeds_all[0]:
all_prompt_embeds = [embed["prompt_embeds"] for embed in prompt_embeds_all]
diff --git a/simpletuner/helpers/caching/vae.py b/simpletuner/helpers/caching/vae.py
index 1fb57a4d4..330568200 100644
--- a/simpletuner/helpers/caching/vae.py
+++ b/simpletuner/helpers/caching/vae.py
@@ -1059,6 +1059,7 @@ def prepare_video_latents(self, samples):
"kandinsky5_video",
"hunyuanvideo",
"longcat_video",
+ "minimaxh3",
]:
if samples.ndim == 4:
original_shape = samples.shape
diff --git a/simpletuner/helpers/configuration/cmd_args.py b/simpletuner/helpers/configuration/cmd_args.py
index 1505d5c1e..65bd66af5 100644
--- a/simpletuner/helpers/configuration/cmd_args.py
+++ b/simpletuner/helpers/configuration/cmd_args.py
@@ -29,6 +29,7 @@
from simpletuner.helpers.training.multi_process import should_log
from simpletuner.helpers.training.optimizer_param import is_optimizer_deprecated, is_optimizer_grad_fp32
from simpletuner.helpers.training.quantisation import MANUAL_QUANTIZATION_PRESETS, PIPELINE_QUANTIZATION_PRESETS
+from simpletuner.helpers.training.sdnq_compile import configure_sdnq_compile_mode
from simpletuner.helpers.training.state_tracker import StateTracker
from simpletuner.simpletuner_sdk.server.services.field_registry.types import (
ConfigField,
@@ -806,6 +807,8 @@ def _normalize_input_args(raw_args):
else:
raise ValueError(f"{option_name} must be a JSON object or a path to a JSON object.")
+ configure_sdnq_compile_mode(getattr(args, "sdnq_compile_mode", "auto"))
+
manual_quant_precisions = set(MANUAL_QUANTIZATION_PRESETS)
pipeline_quant_precisions = set(PIPELINE_QUANTIZATION_PRESETS)
manual_only_precisions = manual_quant_precisions - pipeline_quant_precisions
diff --git a/simpletuner/helpers/data_backend/base.py b/simpletuner/helpers/data_backend/base.py
index 4fc566735..3a73363c9 100644
--- a/simpletuner/helpers/data_backend/base.py
+++ b/simpletuner/helpers/data_backend/base.py
@@ -136,6 +136,7 @@ def _decompress_torch(self, gzip_data: BytesIO):
decompressed_data = file.read()
except Exception as e:
# Handle decompression errors
+ gzip_data.seek(0)
return gzip_data
return BytesIO(decompressed_data)
diff --git a/simpletuner/helpers/data_backend/builders/webshart.py b/simpletuner/helpers/data_backend/builders/webshart.py
index 7a7ec57b5..c9088673c 100644
--- a/simpletuner/helpers/data_backend/builders/webshart.py
+++ b/simpletuner/helpers/data_backend/builders/webshart.py
@@ -29,6 +29,9 @@ def _create_backend(self, config: BaseBackendConfig) -> WebshartDataBackend:
is_mock_backend = hasattr(backend_cls, "_mock_children")
webshart_config = getattr(config, "webshart", None) or {}
+ shard_cache_gb = getattr(config, "webshart_shard_cache_gb", None)
+ if shard_cache_gb is None:
+ shard_cache_gb = 25.0
cache_dir = getattr(config, "webshart_cache_dir", None)
if not cache_dir:
cache_dir = self._default_cache_dir(config)
@@ -45,7 +48,7 @@ def _create_backend(self, config: BaseBackendConfig) -> WebshartDataBackend:
"cache_dir": str(cache_dir) if cache_dir is not None else None,
"metadata_cache_dir": webshart_config.get("metadata_cache_dir"),
"shard_cache_dir": webshart_config.get("shard_cache_dir"),
- "shard_cache_gb": getattr(config, "webshart_shard_cache_gb", None) or 25.0,
+ "shard_cache_gb": shard_cache_gb,
"parallel_downloads": getattr(config, "webshart_parallel_downloads", None) or 4,
"buffer_size": getattr(config, "webshart_buffer_size", None) or 100,
"max_file_size": getattr(config, "webshart_max_file_size", None) or 500 * 1024 * 1024,
diff --git a/simpletuner/helpers/data_backend/config_rules.py b/simpletuner/helpers/data_backend/config_rules.py
index a27cae297..4c2b8404d 100644
--- a/simpletuner/helpers/data_backend/config_rules.py
+++ b/simpletuner/helpers/data_backend/config_rules.py
@@ -12,9 +12,8 @@
make_required_rule,
)
from simpletuner.helpers.data_backend.dataset_types import DatasetType
-from simpletuner.helpers.distillation.registry import DistillationRegistry
+from simpletuner.helpers.distillation.composition import resolve_configured_distiller_requirement_profile
from simpletuner.helpers.distillation.requirements import (
- EMPTY_PROFILE,
DistillerRequirementProfile,
describe_requirement_groups,
evaluate_requirement_profile,
@@ -207,10 +206,7 @@ def _distillation_method_from_config(config: dict) -> Optional[str]:
def _resolve_distiller_profile(config: dict) -> DistillerRequirementProfile:
- method = _distillation_method_from_config(config)
- if not method:
- return EMPTY_PROFILE
- return DistillationRegistry.get_requirement_profile(method)
+ return resolve_configured_distiller_requirement_profile(config)
def _relaxes_training_requirement(profile: DistillerRequirementProfile) -> bool:
diff --git a/simpletuner/helpers/data_backend/factory.py b/simpletuner/helpers/data_backend/factory.py
index 05cffac64..0a0957270 100644
--- a/simpletuner/helpers/data_backend/factory.py
+++ b/simpletuner/helpers/data_backend/factory.py
@@ -124,7 +124,7 @@ def _coerce_bucket_keys(indices: Dict[Any, Iterable]) -> Dict[Any, list]:
)
from simpletuner.helpers.data_backend.webshart import WebshartDataBackend
from simpletuner.helpers.distillation.common import DistillationBase
-from simpletuner.helpers.distillation.registry import DistillationRegistry
+from simpletuner.helpers.distillation.composition import resolve_configured_distiller_requirement_profile
from simpletuner.helpers.distillation.requirements import (
EMPTY_PROFILE,
DistillerRequirementProfile,
@@ -1736,7 +1736,12 @@ def _caption_batches_supported(self) -> bool:
method = self.distillation_method or getattr(self.args, "distillation_method", None)
if method:
try:
- method_profile = DistillationRegistry.get_requirement_profile(method)
+ method_profile = resolve_configured_distiller_requirement_profile(
+ {
+ "distillation_method": method,
+ "distillation_config": getattr(self.args, "distillation_config", None),
+ }
+ )
if method_profile.requires_dataset_type(DatasetType.CAPTION):
return True
except Exception:
@@ -2307,38 +2312,16 @@ def process_conditioning_datasets(self, data_backend_config: List[Dict[str, Any]
f"instance_data_dir={backend.get('instance_data_dir')}"
)
if conditioning_spec_count == 0 and len(linked_conditioning) == 0:
- virtual_id = f"{backend['id']}_conditioning_i2v"
- if any(cfg.get("id") == virtual_id for cfg in data_backend_config):
- info_log(
- f"(id={backend['id']}) I2V conditioning dataset {virtual_id} already present; skipping regeneration."
- )
- else:
- info_log(
- f"(id={backend['id']}) No explicit conditioning datasets provided; creating virtual I2V conditioning dataset {virtual_id}."
- )
- virtual_backend = deepcopy(backend)
- virtual_backend["id"] = virtual_id
- virtual_backend["dataset_type"] = "conditioning"
- virtual_backend.pop("conditioning", None)
- # Conditioning datasets don't use audio - remove audio settings
- # (e.g., IC-LoRA reference videos are visual-only conditioning)
- virtual_backend.pop("audio", None)
- virtual_backend.pop("s2v_datasets", None)
- virtual_backend.pop("_s2v_audio_autoinjected", None)
- virtual_backend["conditioning_data"] = []
- virtual_backend["conditioning_type"] = "reference_strict"
- virtual_backend["source_dataset_id"] = backend["id"]
- virtual_backend["auto_generated"] = False
- # ensure video stanza exists for downstream size alignment
- if isinstance(virtual_backend.get("video"), dict):
- virtual_backend["video"] = dict(virtual_backend["video"])
- virtual_backend["video"].setdefault("is_i2v", True)
- if backend.get("cache_dir_vae"):
- virtual_backend["cache_dir_vae"] = os.path.join(backend["cache_dir_vae"], virtual_id)
- else:
- virtual_backend["cache_dir_vae"] = os.path.join(self.args.cache_dir, "vae", virtual_id)
- backend.setdefault("conditioning_data", []).append(virtual_id)
- conditioning_datasets.append(virtual_backend)
+ info_log(
+ f"(id={backend['id']}) No explicit conditioning datasets provided; "
+ "creating first-frame I2V conditioning data."
+ )
+ backend["conditioning"] = [
+ {
+ "type": "i2v_first_frame",
+ "conditioning_type": "reference_strict",
+ }
+ ]
conditioning_block = backend.get("conditioning", None)
has_explicit_conditioning = conditioning_block not in (None, [], {})
@@ -2573,7 +2556,9 @@ def configure_text_embed_backends(self, data_backend_config: List[Dict[str, Any]
if text_cache_ondemand:
info_log("Skipping null embedding pre-computation for on-demand text cache.")
elif not should_precompute_dropout:
- info_log("Skipping null embedding pre-computation because caption dropout is disabled.")
+ info_log(
+ "Skipping global null embedding pre-computation because caption dropout is disabled or model-specific."
+ )
else:
info_log("Pre-computing null embedding")
logger.debug(f"rank {get_rank()} may skip computing the embedding..")
@@ -3254,11 +3239,10 @@ def _handle_bucket_operations(
if self._is_multi_process():
self.accelerator.wait_for_everyone()
- # When the main process rebuilds buckets (e.g., after cache deletion), ensure
- # other ranks reload the freshly written cache before splitting buckets.
+ # When the main process rebuilds buckets, make every rank partition the same
+ # serialized cache rather than mixing rank 0's in-memory order with reloaded order.
if (
self._is_multi_process()
- and not self.accelerator.is_main_process
and not backend.get("auto_generated", False)
and "aspect" not in self.args.skip_file_discovery
and "aspect" not in backend.get("skip_file_discovery", "")
@@ -3421,6 +3405,23 @@ def _handle_config_versioning(self, backend: Dict[str, Any], init_backend: Dict[
)
init_backend["config"][key] = prev_config[key]
+ runtime_linkage_keys = (
+ "conditioning_data",
+ "conditioning",
+ "video",
+ "s2v_datasets",
+ "_s2v_audio_autoinjected",
+ )
+ for key in runtime_linkage_keys:
+ if key in backend:
+ init_backend["config"][key] = backend[key]
+ if isinstance(getattr(init_backend["metadata_backend"], "config", None), dict):
+ init_backend["metadata_backend"].config[key] = backend[key]
+ else:
+ init_backend["config"].pop(key, None)
+ if isinstance(getattr(init_backend["metadata_backend"], "config", None), dict):
+ init_backend["metadata_backend"].config.pop(key, None)
+
# For Hugging Face datasets, always honor the active caption_strategy (e.g., switching from textfile/filename).
if backend.get("type") == "huggingface":
desired_caption_strategy = backend.get("caption_strategy") or init_backend["config"].get("caption_strategy")
@@ -3733,6 +3734,7 @@ def _process_text_embeddings(
if callable(key_builder):
key_value = key_builder(prompt=caption, default_key=key_value, metadata=metadata)
prompt_records.append({"prompt": caption, "key": key_value, "metadata": metadata})
+ self._append_image_context_dropout_prompt_record(prompt_records, key_value, metadata)
# Add entity label records for images with grounding annotations
grounding_label_count = 0
@@ -3760,7 +3762,10 @@ def _process_text_embeddings(
)
init_backend["text_embed_cache"].compute_embeddings_for_prompts(
- prompt_records, return_concat=False, load_from_cache=False
+ prompt_records,
+ return_concat=False,
+ load_from_cache=False,
+ split_between_processes=False,
)
info_log(f"(id={init_backend['id']}) Completed processing {len(captions)} captions.")
@@ -3869,22 +3874,58 @@ def _process_deferred_text_embeddings(self) -> None:
metadata["image_path"] = image_paths[0]
metadata["data_backend_id"] = data_backend_ids[0]
+ metadata_builder = getattr(self.model, "text_embed_cache_metadata_for_filepath", None)
+ if callable(metadata_builder):
+ metadata.update(
+ metadata_builder(
+ init_backend=init_backend,
+ image_path=image_path_str,
+ prompt=caption,
+ data_backend_id=dataset_id,
+ dataset_relative_path=normalized_identifier,
+ )
+ )
+
if key_type is TextEmbedCacheKey.DATASET_AND_FILENAME:
key_value = f"{dataset_id}:{normalized_identifier}"
elif key_type is TextEmbedCacheKey.FILENAME:
key_value = normalize_data_path(image_path_str, None)
else:
key_value = caption
+ key_builder = getattr(self.model, "text_embed_cache_key_value", None)
+ if callable(key_builder):
+ key_value = key_builder(prompt=caption, default_key=key_value, metadata=metadata)
prompt_records.append({"prompt": caption, "key": key_value, "metadata": metadata})
+ self._append_image_context_dropout_prompt_record(prompt_records, key_value, metadata)
init_backend["text_embed_cache"].compute_embeddings_for_prompts(
- prompt_records, return_concat=False, load_from_cache=False
+ prompt_records,
+ return_concat=False,
+ load_from_cache=False,
+ split_between_processes=False,
)
info_log(f"(id={dataset_id}) Completed processing {len(captions)} captions with image context.")
# Clear the deferred queue
self._deferred_text_embed_backends.clear()
+ def _append_image_context_dropout_prompt_record(
+ self,
+ prompt_records: List[Dict[str, Any]],
+ default_key: str,
+ metadata: Dict[str, Any],
+ ) -> None:
+ if getattr(self.args, "caption_dropout_probability", 0.1) <= 0.0:
+ return
+ if not getattr(self.model, "uses_image_context_dropout_caption_cache", lambda: False)():
+ return
+ dropout_metadata = dict(metadata)
+ key_builder = getattr(self.model, "text_embed_cache_key_value", None)
+ key_value = default_key
+ if callable(key_builder):
+ key_value = key_builder(prompt="", default_key=default_key, metadata=dropout_metadata)
+ prompt_records.append({"prompt": "", "key": key_value, "metadata": dropout_metadata})
+
def _handle_auto_generated_dataset(self, backend: Dict[str, Any], init_backend: Dict[str, Any]) -> None:
"""Handle auto-generated reference datasets."""
# we have to auto-generate the data for the reference images.
@@ -4049,7 +4090,7 @@ def _configure_vae_cache(
if dataset_type_enum is DatasetType.AUDIO:
uses_audio_latents = False
try:
- uses_audio_latents = bool(self.model.uses_audio_latents())
+ uses_audio_latents = bool(self.model.uses_audio_latents_for_data_backend(init_backend.get("id")))
except AttributeError:
uses_audio_latents = False
if not uses_audio_latents:
diff --git a/simpletuner/helpers/data_backend/webshart.py b/simpletuner/helpers/data_backend/webshart.py
index 3c16c4df0..2ec5ec3ba 100644
--- a/simpletuner/helpers/data_backend/webshart.py
+++ b/simpletuner/helpers/data_backend/webshart.py
@@ -1,16 +1,20 @@
import json
import logging
import os
+import random
+import time
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Any, List, Optional, Tuple, Union
+import requests
import torch
from simpletuner.helpers.data_backend.base import BaseDataBackend
from simpletuner.helpers.data_backend.dataset_types import DatasetType, ensure_dataset_type
-from simpletuner.helpers.image_manipulation.load import load_image
+from simpletuner.helpers.image_manipulation.load import load_image, load_video
+from simpletuner.helpers.training import video_file_extensions
from simpletuner.helpers.training.multi_process import should_log
logger = logging.getLogger("WebshartDataBackend")
@@ -29,6 +33,7 @@ class WebshartSampleRef:
class WebshartDataBackend(BaseDataBackend):
SAMPLE_PREFIX = "webshart://"
+ PATH_NORMALIZED_SAMPLE_PREFIX = "webshart:/"
CACHE_EXTENSIONS = {".json", ".pt", ".msgpack", ".safetensors"}
def __init__(
@@ -69,8 +74,12 @@ def __init__(
self.metadata_cache_dir = (
str(metadata_cache_dir) if metadata_cache_dir else str(Path(self.cache_dir) / "metadata_cache")
)
- self.shard_cache_dir = str(shard_cache_dir) if shard_cache_dir else str(Path(self.cache_dir) / "shard_cache")
self.shard_cache_gb = float(shard_cache_gb)
+ if self.shard_cache_gb < 0:
+ raise ValueError("shard_cache_gb must be non-negative; use 0 to disable whole-shard caching.")
+ self.shard_cache_dir = str(shard_cache_dir) if shard_cache_dir else str(Path(self.cache_dir) / "shard_cache")
+ if self.shard_cache_gb == 0:
+ self.shard_cache_dir = None
self.parallel_downloads = int(parallel_downloads)
self.buffer_size = int(buffer_size)
self.max_file_size = int(max_file_size)
@@ -88,7 +97,7 @@ def __init__(
metadata=self.metadata,
)
self.dataset.enable_metadata_cache(location=self.metadata_cache_dir)
- if self.shard_cache_dir:
+ if self.shard_cache_dir is not None:
Path(self.shard_cache_dir).mkdir(parents=True, exist_ok=True)
self.dataset.enable_shard_cache(
location=self.shard_cache_dir,
@@ -117,6 +126,10 @@ def normalize_sample_id(cls, identifier: Union[str, Path]) -> str:
marker = value.find(cls.SAMPLE_PREFIX)
if marker >= 0:
return value[marker:]
+ marker = value.find(cls.PATH_NORMALIZED_SAMPLE_PREFIX)
+ if marker >= 0:
+ remainder = value[marker + len(cls.PATH_NORMALIZED_SAMPLE_PREFIX) :]
+ return f"{cls.SAMPLE_PREFIX}{remainder}"
return value
@classmethod
@@ -137,10 +150,9 @@ def parse_sample_id(cls, identifier: Union[str, Path]) -> WebshartSampleRef:
@classmethod
def is_sample_id(cls, identifier: Union[str, Path]) -> bool:
- value = str(identifier)
- if cls.SAMPLE_PREFIX not in value:
+ sample_id = cls.normalize_sample_id(identifier)
+ if not sample_id.startswith(cls.SAMPLE_PREFIX):
return False
- sample_id = cls.normalize_sample_id(value)
filename = sample_id.split("/", 2)[-1] if "/" in sample_id else sample_id
return Path(filename).suffix.lower() not in {".pt", ".safetensors"}
@@ -188,8 +200,92 @@ def _is_cache_identifier(self, identifier: Union[str, Path]) -> bool:
def _read_sample_bytes(self, identifier: Union[str, Path]) -> bytes:
sample_ref = self.parse_sample_id(identifier)
- entry = self.loader.load_sample(sample_ref.shard_idx, sample_ref.sample_idx)
- return bytes(entry.data)
+ max_attempts = 6
+ for attempt in range(max_attempts):
+ try:
+ entry = self.loader.load_sample(sample_ref.shard_idx, sample_ref.sample_idx)
+ return bytes(entry.data)
+ except Exception as exc:
+ message = str(exc).lower()
+ retryable = any(
+ marker in message
+ for marker in (
+ "rate limit",
+ "http 429",
+ "status code 429",
+ "connection reset",
+ "temporarily unavailable",
+ "timed out",
+ "timeout",
+ )
+ )
+ if not retryable or attempt + 1 >= max_attempts:
+ raise
+ delay = min(30.0, 2.0**attempt) + random.uniform(0.0, 1.0)
+ logger.warning(
+ "Transient error reading Webshart sample %s; retrying in %.1fs (%d/%d): %s",
+ identifier,
+ delay,
+ attempt + 1,
+ max_attempts,
+ exc,
+ )
+ time.sleep(delay)
+
+ def read_sample_head_tail(
+ self,
+ identifier: Union[str, Path],
+ *,
+ file_metadata: Optional[dict] = None,
+ head_bytes: int = 4096,
+ tail_bytes: int = 131072,
+ ) -> tuple[bytes, bytes, int]:
+ """Range-read the beginning and end of a sample without loading its full TAR member."""
+ sample_ref = self.parse_sample_id(identifier)
+ metadata = file_metadata or self.get_shard_metadata(sample_ref.shard_idx).get(sample_ref.filename, {})
+ offset = metadata.get("offset")
+ length = metadata.get("length", metadata.get("size"))
+ shard_info = self.dataset.get_shard_info(sample_ref.shard_idx)
+ tar_path = shard_info.get("tar_path") if isinstance(shard_info, dict) else None
+ if offset is None or length is None or not str(tar_path or "").startswith(("http://", "https://")):
+ raise ValueError(f"Range metadata is unavailable for Webshart sample {identifier}.")
+
+ offset = int(offset)
+ length = int(length)
+ if length <= 0:
+ raise ValueError(f"Invalid Webshart sample length for {identifier}: {length}")
+
+ token = self.hf_token
+ if token is True:
+ from huggingface_hub import get_token
+
+ token = get_token()
+ base_headers = {"Authorization": f"Bearer {token}"} if token else {}
+
+ def _read_range(relative_start: int, relative_end: int) -> bytes:
+ absolute_start = offset + relative_start
+ absolute_end = offset + relative_end
+ headers = {**base_headers, "Range": f"bytes={absolute_start}-{absolute_end}"}
+ response = requests.get(str(tar_path), headers=headers, stream=True, timeout=(10, 60))
+ try:
+ if response.status_code != 206:
+ raise IOError(f"Range request for {identifier} returned HTTP {response.status_code} instead of 206.")
+ payload = response.content
+ finally:
+ response.close()
+ expected_size = relative_end - relative_start + 1
+ if len(payload) != expected_size:
+ raise IOError(f"Range request for {identifier} returned {len(payload)} bytes; expected {expected_size}.")
+ return payload
+
+ head_size = min(max(1, int(head_bytes)), length)
+ tail_size = min(max(1, int(tail_bytes)), length)
+ head = _read_range(0, head_size - 1)
+ if tail_size == length:
+ tail = head if head_size == length else _read_range(0, length - 1)
+ else:
+ tail = _read_range(length - tail_size, length - 1)
+ return head, tail, length
def _sample_index_for_filename(self, shard_idx: int, filename: str) -> Optional[int]:
shard_idx = int(shard_idx)
@@ -205,6 +301,11 @@ def get_caption(self, image_path: str) -> Optional[str]:
return None
sample_ref = self.parse_sample_id(image_path)
+ sample_metadata = self.get_shard_metadata(sample_ref.shard_idx).get(sample_ref.filename, {}) or {}
+ caption = sample_metadata.get("captions")
+ if caption:
+ return str(caption).strip()
+
caption_filename = Path(sample_ref.filename).with_suffix(".txt").name
caption_sample_idx = self._sample_index_for_filename(sample_ref.shard_idx, caption_filename)
if caption_sample_idx is None:
@@ -296,7 +397,9 @@ def get_abs_path(self, sample_path: str = None) -> Optional[str]:
def read_image(self, filepath: str, delete_problematic_images: bool = False):
try:
- return load_image(self.read(filepath, as_byteIO=True))
+ file_extension = Path(self.normalize_sample_id(filepath)).suffix.lower().strip(".")
+ loader = load_video if file_extension in video_file_extensions else load_image
+ return loader(self.read(filepath, as_byteIO=True))
except Exception as exc:
logger.error("Error opening webshart sample %s: %s", filepath, exc)
if delete_problematic_images:
@@ -322,6 +425,7 @@ def torch_load(self, filename):
data = self.read(filename, as_byteIO=True)
if self.compress_cache:
data = self._decompress_torch(data)
+ data.seek(0)
return torch.load(data, map_location="cpu")
def torch_save(self, data, filename):
diff --git a/simpletuner/helpers/data_generation/conditioning.py b/simpletuner/helpers/data_generation/conditioning.py
index 5a1da547f..0adada714 100644
--- a/simpletuner/helpers/data_generation/conditioning.py
+++ b/simpletuner/helpers/data_generation/conditioning.py
@@ -408,6 +408,9 @@ def discover_all_files(self) -> List[str]:
def generate_target_filename(self, source_filepath: str) -> Tuple[str, str]:
"""map source path to target path, preserving structure"""
base, ext = os.path.splitext(os.path.basename(source_filepath))
+ target_extension = getattr(self.sample_generator, "target_extension", None)
+ if target_extension:
+ ext = target_extension if str(target_extension).startswith(".") else f".{target_extension}"
filename = f"{base}{ext}"
subpath = ""
@@ -488,6 +491,7 @@ def process_buckets(self):
"errors": 0,
"total": len(files),
}
+ dispatched = 0
last_report = 0
for fp in tqdm(files, desc=f"Bucket {bucket}", position=self.rank, leave=False):
@@ -506,6 +510,7 @@ def process_buckets(self):
if self.gpu_mode and sent:
stats["processed"] += sent
elif sent:
+ dispatched += sent
self.debug_log(f"Dispatched batch {batch_id} ({sent} items)")
if not self.gpu_mode:
@@ -537,12 +542,22 @@ def process_buckets(self):
futures = self._process_futures(futures, executor)
while self.process_queue.qsize():
batch_id += 1
- self._process_images_in_batch(batch_id)
+ sent = self._process_images_in_batch(batch_id)
+ if self.gpu_mode and sent:
+ stats["processed"] += sent
+ elif sent:
+ dispatched += sent
if not self.gpu_mode:
start = time.time()
- while (self.transform_queue.qsize() or not self.done_queue.empty()) and time.time() - start < 30:
- stats["processed"] += self._check_completion_queue()
+ while stats["processed"] < dispatched and time.time() - start < 300:
+ completed = self._check_completion_queue()
+ if completed:
+ stats["processed"] += completed
time.sleep(0.1)
+ if stats["processed"] < dispatched:
+ missing = dispatched - stats["processed"]
+ stats["errors"] += missing
+ logger.error(f"(id={self.id}) Timed out waiting for {missing} generated conditioning samples.")
msg = f"(id={self.id}) Bucket {bucket} done: {stats}"
if self.rank == 0:
logger.info(msg)
diff --git a/simpletuner/helpers/data_generation/sample_generator.py b/simpletuner/helpers/data_generation/sample_generator.py
index adb40b604..17d485c5b 100644
--- a/simpletuner/helpers/data_generation/sample_generator.py
+++ b/simpletuner/helpers/data_generation/sample_generator.py
@@ -1195,6 +1195,8 @@ class I2VFirstFrameSampleGenerator(SampleGenerator):
self-contained I2V conditioning datasets.
"""
+ target_extension = ".png"
+
def transform_batch(
self,
images: List[Any],
diff --git a/simpletuner/helpers/distillation/anyflow/distiller.py b/simpletuner/helpers/distillation/anyflow/distiller.py
index 1ecc8c3a0..9d1638db1 100644
--- a/simpletuner/helpers/distillation/anyflow/distiller.py
+++ b/simpletuner/helpers/distillation/anyflow/distiller.py
@@ -91,7 +91,7 @@ def prepare_batch(self, batch: Dict[str, Any], model, state) -> Dict[str, Any]:
r_timesteps = self._timesteps_from_sigmas(r_sigmas, timesteps)
r_timesteps = r_timesteps.to(device=batch["timesteps"].device, dtype=batch["timesteps"].dtype)
- target = self._base_flow_target(batch)
+ target = self._base_flow_target(batch, model=model)
if self.config["target_mode"] == "online_teacher":
target = self._online_teacher_average_velocity(
prepared_batch=batch,
@@ -103,7 +103,7 @@ def prepare_batch(self, batch: Dict[str, Any], model, state) -> Dict[str, Any]:
flowmap_key = getattr(model, "FLOWMAP_R_TIMESTEP_BATCH_KEY", self.FLOWMAP_R_TIMESTEP_BATCH_KEY)
batch[flowmap_key] = r_timesteps
batch["anyflow_r_timesteps"] = r_timesteps
- batch["anyflow_timestep_interval"] = batch["timesteps"].to(dtype=r_timesteps.dtype) - r_timesteps
+ batch["anyflow_timestep_interval"] = (batch["timesteps"].to(dtype=r_timesteps.dtype) - r_timesteps).abs()
batch["target"] = target.detach()
batch["flow_target"] = batch["target"]
return batch
@@ -125,9 +125,10 @@ def compute_distill_loss(
}
if r_timesteps is not None:
logs["anyflow_r_timestep"] = float(torch.mean(r_timesteps.float()).detach())
- logs["anyflow_interval"] = float(
- torch.mean((prepared_batch["timesteps"].float() - r_timesteps.float())).detach()
- )
+ interval = prepared_batch.get("anyflow_timestep_interval")
+ if interval is None:
+ interval = (prepared_batch["timesteps"].float() - r_timesteps.float()).abs()
+ logs["anyflow_interval"] = float(torch.mean(interval.float()).detach())
return loss, logs
def get_scheduler(self, scheduler=None):
@@ -238,6 +239,9 @@ def _scalar_sigmas(prepared_batch: Dict[str, Any]) -> torch.Tensor:
return sigmas.reshape(batch_size, -1)[:, 0].clamp(0.0, 1.0)
def _timesteps_from_sigmas(self, sigmas: torch.Tensor, reference_timesteps: torch.Tensor) -> torch.Tensor:
+ converter = getattr(self.teacher_model, "flow_matching_timesteps_from_sigmas", None)
+ if callable(converter):
+ return converter(sigmas, reference_timesteps=reference_timesteps)
if torch.max(reference_timesteps.detach().float()) <= 1.0:
return sigmas
return sigmas * self.num_train_timesteps
@@ -261,13 +265,29 @@ def _validate_interval(t_sigmas: torch.Tensor, r_sigmas: torch.Tensor) -> None:
"or use a sampler that leaves a positive interval."
)
- @staticmethod
- def _base_flow_target(prepared_batch: Dict[str, Any]) -> torch.Tensor:
+ def _base_flow_target(self, prepared_batch: Dict[str, Any], model=None) -> torch.Tensor:
+ target_model = model or self.student_model or self.teacher_model
+ get_target = getattr(target_model, "get_flow_matching_target", None)
+ if callable(get_target):
+ target = get_target(prepared_batch, prefer_explicit_target=True)
+ return target.to(device=prepared_batch["latents"].device, dtype=prepared_batch["latents"].dtype)
flow_target = prepared_batch.get("flow_target")
if torch.is_tensor(flow_target):
return flow_target.to(device=prepared_batch["latents"].device, dtype=prepared_batch["latents"].dtype)
return prepared_batch["noise"] - prepared_batch["latents"]
+ def _teacher_prediction_to_noiseward_flow(self, prediction: torch.Tensor) -> torch.Tensor:
+ converter = getattr(self.teacher_model, "prediction_to_noiseward_flow", None)
+ if callable(converter):
+ return converter(prediction)
+ return prediction
+
+ def _noiseward_flow_to_student_prediction(self, flow: torch.Tensor) -> torch.Tensor:
+ converter = getattr(self.student_model, "noiseward_flow_to_prediction", None)
+ if callable(converter):
+ return converter(flow)
+ return flow
+
def _online_teacher_average_velocity(
self,
*,
@@ -287,15 +307,17 @@ def _online_teacher_average_velocity(
teacher_batch = self._teacher_batch(prepared_batch, current_latents, current_sigmas)
teacher_prediction = self.teacher_model.model_predict(teacher_batch)["model_prediction"]
teacher_prediction = teacher_prediction.to(device=current_latents.device, dtype=current_latents.dtype)
+ noiseward_prediction = self._teacher_prediction_to_noiseward_flow(teacher_prediction)
step = self._broadcast_time(next_sigmas - current_sigmas, current_latents)
- current_latents = current_latents + step * teacher_prediction
+ current_latents = current_latents + step * noiseward_prediction
current_sigmas = next_sigmas
finally:
self.toggle_adapter(enable=True)
denominator = self._broadcast_time(r_sigmas - t_sigmas, current_latents)
average_velocity = (current_latents - start_latents) / denominator
- return average_velocity.to(device=base_target.device, dtype=base_target.dtype)
+ target = self._noiseward_flow_to_student_prediction(average_velocity)
+ return target.to(device=base_target.device, dtype=base_target.dtype)
def _rollout_sigma_schedule(self, t_sigmas: torch.Tensor, r_sigmas: torch.Tensor):
steps = int(self.config["teacher_rollout_steps"])
diff --git a/simpletuner/helpers/distillation/anyflow/scheduler.py b/simpletuner/helpers/distillation/anyflow/scheduler.py
index 54946c22c..8f0257b44 100644
--- a/simpletuner/helpers/distillation/anyflow/scheduler.py
+++ b/simpletuner/helpers/distillation/anyflow/scheduler.py
@@ -151,14 +151,6 @@ def _scheduler_timesteps(self, device: torch.device) -> torch.Tensor:
return timesteps.detach().to(device=device, dtype=torch.float32).reshape(-1)
return torch.tensor(list(timesteps), device=device, dtype=torch.float32).reshape(-1)
- def _scheduler_sigmas(self, device: torch.device) -> Optional[torch.Tensor]:
- sigmas = getattr(self.scheduler, "sigmas", None)
- if sigmas is None:
- return None
- if torch.is_tensor(sigmas):
- return sigmas.detach().to(device=device, dtype=torch.float32).reshape(-1)
- return torch.tensor(list(sigmas), device=device, dtype=torch.float32).reshape(-1)
-
def _train_timestep_scale(self, schedule: torch.Tensor) -> float:
if self.num_train_timesteps is not None:
return float(self.num_train_timesteps)
@@ -176,13 +168,8 @@ def _endpoint_schedule(
train_scale: float,
device: torch.device,
) -> torch.Tensor:
- sigmas = self._scheduler_sigmas(device)
- if sigmas is not None and sigmas.numel() == schedule.numel() + 1:
- sigma_max = torch.max(torch.abs(sigmas)).item()
- endpoint_schedule = sigmas[1:] * train_scale if sigma_max <= 1.5 else sigmas[1:]
- return endpoint_schedule.to(device=device, dtype=torch.float32)
-
- final_raw_timestep = 0.0 if schedule[0] >= schedule[-1] else train_scale
+ schedule_scale = 1.0 if torch.max(torch.abs(schedule)).item() <= 1.5 else train_scale
+ final_raw_timestep = 0.0 if schedule[0] >= schedule[-1] else schedule_scale
final = torch.tensor([final_raw_timestep], device=device, dtype=torch.float32)
return torch.cat([schedule[1:], final])
diff --git a/simpletuner/helpers/distillation/composition.py b/simpletuner/helpers/distillation/composition.py
new file mode 100644
index 000000000..ba77fa757
--- /dev/null
+++ b/simpletuner/helpers/distillation/composition.py
@@ -0,0 +1,60 @@
+"""Helpers for resolving composed distillation configurations."""
+
+from __future__ import annotations
+
+from typing import Any, Mapping, Optional
+
+import simpletuner.helpers.distillation.factory # noqa: F401
+from simpletuner.helpers.distillation.registry import DistillationRegistry
+from simpletuner.helpers.distillation.requirements import (
+ EMPTY_PROFILE,
+ DistillerRequirementProfile,
+ merge_distiller_requirement_profiles,
+)
+
+
+def _config_value(config: Any, key: str, default: Any = None) -> Any:
+ if isinstance(config, Mapping):
+ return config.get(key, default)
+ return getattr(config, key, default)
+
+
+def _normalize_method(value: Any) -> Optional[str]:
+ if value in (None, "", False):
+ return None
+ method = str(value).strip().lower().replace("-", "_")
+ if method in {"none", "false", "0"}:
+ return None
+ return method
+
+
+def _method_config(method: str, distillation_config: Any) -> Mapping[str, Any]:
+ if not isinstance(distillation_config, Mapping):
+ return {}
+ specific = distillation_config.get(method)
+ if isinstance(specific, Mapping):
+ return specific
+ return distillation_config
+
+
+def resolve_configured_distiller_requirement_profile(config: Any) -> DistillerRequirementProfile:
+ """Resolve top-level and H3 drift inner-distiller data requirements from a config object."""
+ method = _normalize_method(_config_value(config, "distillation_method"))
+ if method is None:
+ method = _normalize_method(_config_value(config, "--distillation_method"))
+ if method is None:
+ return EMPTY_PROFILE
+
+ profile = DistillationRegistry.get_requirement_profile(method)
+ if method != "h3_drift":
+ return profile
+
+ h3_config = _method_config(method, _config_value(config, "distillation_config"))
+ inner_method = _normalize_method(h3_config.get("inner_distillation_method"))
+ if inner_method is None:
+ return profile
+ if inner_method == "h3_drift":
+ raise ValueError("H3 drift may not wrap another h3_drift distiller.")
+
+ inner_profile = DistillationRegistry.get_requirement_profile(inner_method)
+ return merge_distiller_requirement_profiles(profile, inner_profile)
diff --git a/simpletuner/helpers/distillation/dmd/distiller.py b/simpletuner/helpers/distillation/dmd/distiller.py
index a82a3302a..5dbfefe68 100644
--- a/simpletuner/helpers/distillation/dmd/distiller.py
+++ b/simpletuner/helpers/distillation/dmd/distiller.py
@@ -85,6 +85,7 @@ def __init__(
self.fake_score_transformer,
self.scheduler_adapter,
self.weight_dtype,
+ foundation=self.teacher_model,
)
self.denoising_steps = self._parse_denoising_steps(self.config["dmd_denoising_steps"])
diff --git a/simpletuner/helpers/distillation/factory.py b/simpletuner/helpers/distillation/factory.py
index e64b7bdab..9d5e42822 100644
--- a/simpletuner/helpers/distillation/factory.py
+++ b/simpletuner/helpers/distillation/factory.py
@@ -6,8 +6,9 @@
# Ensure registry-backed distillers (like self_forcing) register themselves on import.
import simpletuner.helpers.distillation.anyflow # noqa: F401
import simpletuner.helpers.distillation.flow_dpo # noqa: F401
+import simpletuner.helpers.distillation.h3_drift # noqa: F401
import simpletuner.helpers.distillation.perflow.distiller # noqa: F401
-import simpletuner.helpers.distillation.self_forcing # noqa: F401
+import simpletuner.helpers.distillation.self_forcing.distiller # noqa: F401
from simpletuner.helpers.distillation.common import DistillationBase, validate_distillation_text_encoder_training
from simpletuner.helpers.distillation.registry import DistillationRegistry
@@ -23,6 +24,7 @@ class DistillationMethod(Enum):
PERFLOW = "perflow"
FLOW_DPO = "flow_dpo"
ANYFLOW = "anyflow"
+ H3_DRIFT = "h3_drift"
SELF_FORCING = "self_forcing"
@classmethod
@@ -147,6 +149,19 @@ def create_distiller(
},
student_model=student_model,
)
+ elif method == DistillationMethod.H3_DRIFT:
+ return DistillerFactory._create_registered_distiller(
+ registry_key=method.value,
+ teacher_model=teacher_model,
+ noise_scheduler=noise_scheduler,
+ distill_config=distill_config,
+ runtime_config_defaults={
+ "model_type": model_type,
+ "model_family": model_family,
+ "prediction_type": prediction_type,
+ },
+ student_model=student_model,
+ )
elif method == DistillationMethod.SELF_FORCING:
return DistillerFactory._create_registered_distiller(
registry_key=method.value,
diff --git a/simpletuner/helpers/distillation/flow_dpo/distiller.py b/simpletuner/helpers/distillation/flow_dpo/distiller.py
index a95cad124..814d28c63 100644
--- a/simpletuner/helpers/distillation/flow_dpo/distiller.py
+++ b/simpletuner/helpers/distillation/flow_dpo/distiller.py
@@ -85,8 +85,8 @@ def compute_distill_loss(
finally:
self.toggle_adapter(enable=True)
- win_target = prepared_batch["noise"] - win_latents
- lose_target = prepared_batch["noise"] - lose_latents
+ win_target = self._flow_target(prepared_batch, win_latents)
+ lose_target = self._flow_target(lose_batch, lose_latents)
mask = self._mask_for_loss(prepared_batch, policy_win)
policy_win_err = self._per_sample_error(policy_win, win_target, mask)
@@ -175,6 +175,13 @@ def _build_rejected_batch(prepared_batch: Dict[str, Any], rejected_latents: torc
] * prepared_batch["input_noise"]
return rejected_batch
+ def _flow_target(self, prepared_batch: Dict[str, Any], latents: torch.Tensor) -> torch.Tensor:
+ get_target = getattr(self.student_model, "get_flow_matching_target", None)
+ if callable(get_target):
+ target = get_target(prepared_batch, latents=latents, prefer_explicit_target=False)
+ return target.to(device=latents.device, dtype=latents.dtype)
+ return prepared_batch["noise"] - latents
+
def _mask_for_loss(self, prepared_batch: Dict[str, Any], prediction: torch.Tensor) -> Optional[torch.Tensor]:
loss_mask_type = prepared_batch.get("loss_mask_type")
if loss_mask_type not in {"mask", "segmentation"}:
diff --git a/simpletuner/helpers/distillation/h3_drift/__init__.py b/simpletuner/helpers/distillation/h3_drift/__init__.py
new file mode 100644
index 000000000..1c9745af0
--- /dev/null
+++ b/simpletuner/helpers/distillation/h3_drift/__init__.py
@@ -0,0 +1,3 @@
+from .distiller import H3DriftDistiller
+
+__all__ = ["H3DriftDistiller"]
diff --git a/simpletuner/helpers/distillation/h3_drift/distiller.py b/simpletuner/helpers/distillation/h3_drift/distiller.py
new file mode 100644
index 000000000..cf8e89d28
--- /dev/null
+++ b/simpletuner/helpers/distillation/h3_drift/distiller.py
@@ -0,0 +1,513 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn.functional as F
+
+from simpletuner.helpers.distillation.common import DistillationBase
+from simpletuner.helpers.distillation.registry import DistillationRegistry
+
+
+@dataclass(frozen=True)
+class _H3Prediction:
+ video: Optional[torch.Tensor]
+ audio: Optional[torch.Tensor]
+
+
+@dataclass(frozen=True)
+class _H3JointLoss:
+ loss: torch.Tensor
+ video_loss: torch.Tensor
+ audio_loss: torch.Tensor
+ video_elements: int
+ audio_elements: int
+
+
+class H3DriftDistiller(DistillationBase):
+ """Regularize MiniMax-H3 LoRA/LyCORIS training against the frozen base prediction."""
+
+ _FLOWMAP_BATCH_KEYS = (
+ "flowmap_r_timesteps",
+ "anyflow_r_timesteps",
+ "anyflow_timestep_interval",
+ )
+ _DEFAULTS: Dict[str, Any] = {
+ "distillation_type": "h3_drift",
+ "loss_weight": 1.0,
+ "sft_loss_weight": 1.0,
+ "balance": "token",
+ "video_weight": 1.0,
+ "audio_weight": 1.0,
+ "inner_distillation_method": None,
+ "inner_distillation_config": {},
+ }
+
+ def __init__(
+ self,
+ teacher_model,
+ student_model=None,
+ *,
+ noise_scheduler=None,
+ config: Optional[Dict[str, Any]] = None,
+ ):
+ merged_config = dict(self._DEFAULTS)
+ if config:
+ merged_config.update(config)
+
+ super().__init__(teacher_model, student_model, merged_config)
+ self.noise_scheduler = noise_scheduler
+
+ model_family = str(self.config.get("model_family") or "").lower().replace("_", "")
+ model_name = str(getattr(teacher_model, "NAME", "") or "").lower().replace(" ", "").replace("-", "")
+ if model_family not in {"minimaxh3", ""} and model_name != "minimaxh3":
+ raise ValueError("H3 drift distillation only supports MiniMax-H3.")
+ if not self.is_flow_matching:
+ raise ValueError("H3 drift distillation requires a flow-matching model.")
+ if not self.low_rank_distillation or self.config.get("model_type") != "lora":
+ raise ValueError("H3 drift distillation only supports low-rank LoRA/LyCORIS training.")
+
+ balance = str(self.config.get("balance", "token")).lower()
+ if balance not in {"token", "modality"}:
+ raise ValueError("H3 drift balance must be one of: token, modality.")
+ self.config["balance"] = balance
+
+ video_weight = float(self.config.get("video_weight", 1.0))
+ audio_weight = float(self.config.get("audio_weight", 1.0))
+ if video_weight < 0 or audio_weight < 0 or video_weight + audio_weight <= 0:
+ raise ValueError("H3 drift video/audio weights must be non-negative and not both zero.")
+ self.config["video_weight"] = video_weight
+ self.config["audio_weight"] = audio_weight
+ self.inner_distiller = self._create_inner_distiller()
+
+ def _create_inner_distiller(self):
+ raw_method = self.config.get("inner_distillation_method")
+ if raw_method in (None, "", False):
+ return None
+
+ method = str(raw_method).strip().lower()
+ if method in {"none", "false", "0"}:
+ return None
+ if method == "h3_drift":
+ raise ValueError("H3 drift may not wrap another h3_drift distiller.")
+
+ inner_config = self.config.get("inner_distillation_config") or {}
+ if not isinstance(inner_config, dict):
+ raise ValueError("H3 drift inner_distillation_config must be a mapping.")
+
+ from simpletuner.helpers.distillation.factory import DistillerFactory
+
+ return DistillerFactory.create_distiller(
+ method,
+ teacher_model=self.teacher_model,
+ noise_scheduler=self.noise_scheduler,
+ config={
+ "distillation_config": {method: inner_config},
+ "train_text_encoder": bool(self.config.get("train_text_encoder", False)),
+ },
+ model_type=str(self.config.get("model_type") or "lora"),
+ model_family=self.config.get("model_family"),
+ prediction_type=self.config.get("prediction_type"),
+ student_model=None if self.low_rank_distillation else self.student_model,
+ )
+
+ def requires_distillation_cache(self) -> bool:
+ return bool(self.inner_distiller and self.inner_distiller.requires_distillation_cache())
+
+ def get_required_distillation_cache_type(self) -> Optional[str]:
+ if self.inner_distiller is None:
+ return None
+ return self.inner_distiller.get_required_distillation_cache_type()
+
+ def get_ode_generator_provider(self):
+ if self.inner_distiller is None:
+ return None
+ return self.inner_distiller.get_ode_generator_provider()
+
+ def get_scheduler(self, scheduler=None):
+ if self.inner_distiller is None:
+ return super().get_scheduler(scheduler)
+ return self.inner_distiller.get_scheduler(scheduler)
+
+ def prepare_batch(self, batch, model, state):
+ if self.inner_distiller is None:
+ return batch
+ return self.inner_distiller.prepare_batch(batch, model, state)
+
+ def consumes_caption_batches(self) -> bool:
+ return bool(self.inner_distiller and self.inner_distiller.consumes_caption_batches())
+
+ def prepare_caption_batch(self, caption_batch: Dict[str, Any], model, state) -> Dict[str, Any]:
+ if self.inner_distiller is None:
+ return super().prepare_caption_batch(caption_batch, model, state)
+ return self.inner_distiller.prepare_caption_batch(caption_batch, model, state)
+
+ def pre_training_step(self, model, step):
+ if self.inner_distiller is not None:
+ self.inner_distiller.pre_training_step(model, step)
+
+ def post_training_step(self, model, step):
+ if self.inner_distiller is not None:
+ self.inner_distiller.post_training_step(model, step)
+
+ def compute_distill_loss(
+ self,
+ prepared_batch: Dict[str, Any],
+ model_output: Dict[str, Any],
+ original_loss: torch.Tensor,
+ ):
+ drift_loss_weight = float(self.config.get("loss_weight", 1.0))
+ if drift_loss_weight == 0.0:
+ zero = original_loss.detach() * 0.0
+ joint_loss = _H3JointLoss(
+ loss=zero,
+ video_loss=zero,
+ audio_loss=zero,
+ video_elements=0,
+ audio_elements=0,
+ )
+ else:
+ prediction = self._prediction_from_output(model_output)
+ try:
+ self.toggle_adapter(enable=False)
+ with torch.no_grad():
+ reference_output = self.teacher_model.model_predict(prepared_batch)
+ finally:
+ self.toggle_adapter(enable=True)
+
+ reference = self._prediction_from_output(reference_output)
+ self._clear_reference_buffers(reference_output)
+
+ joint_loss = self._joint_prediction_loss(
+ prediction,
+ reference,
+ video_mask=self._video_mask_for_loss(prepared_batch, prediction.video),
+ audio_mask=prepared_batch.get("audio_latent_mask"),
+ sample_weight=prepared_batch.get("sample_weight"),
+ balance=self.config["balance"],
+ video_weight=self.config["video_weight"],
+ audio_weight=self.config["audio_weight"],
+ )
+
+ sft_loss_weight = float(self.config.get("sft_loss_weight", 1.0))
+ logs = {}
+ if self.inner_distiller is not None:
+ current_loss, logs = self.inner_distiller.compute_distill_loss(prepared_batch, model_output, original_loss)
+ logs = self._inner_logs(logs)
+ sft_loss = self._normal_h3_loss(prepared_batch, model_output, original_loss) * sft_loss_weight
+ current_loss = current_loss + sft_loss
+ else:
+ sft_loss = original_loss * sft_loss_weight
+ current_loss = sft_loss
+
+ drift_loss = joint_loss.loss * drift_loss_weight
+ loss = current_loss + drift_loss
+
+ logs.update(
+ {
+ "h3_drift_loss": float(joint_loss.loss.detach()),
+ "h3_drift_video_loss": float(joint_loss.video_loss.detach()),
+ "h3_drift_audio_loss": float(joint_loss.audio_loss.detach()),
+ "h3_drift_video_elements": float(joint_loss.video_elements),
+ "h3_drift_audio_elements": float(joint_loss.audio_elements),
+ "h3_drift_weighted_loss": float(drift_loss.detach()),
+ "total": float(loss.detach()),
+ }
+ )
+ if sft_loss_weight != 0.0:
+ logs["h3_drift_sft_loss"] = float(sft_loss.detach())
+ return loss, logs
+
+ def _normal_h3_loss(
+ self,
+ prepared_batch: Dict[str, Any],
+ model_output: Dict[str, Any],
+ original_loss: torch.Tensor,
+ ) -> torch.Tensor:
+ sft_loss_weight = float(self.config.get("sft_loss_weight", 1.0))
+ if sft_loss_weight == 0.0:
+ return original_loss * 0.0
+ if self.inner_distiller is None:
+ return original_loss
+
+ normal_batch = self._normal_h3_batch(prepared_batch)
+ if self._has_inner_timestep_conditioning(prepared_batch):
+ normal_output = self.teacher_model.model_predict(normal_batch)
+ try:
+ return self.teacher_model.loss(
+ normal_batch,
+ normal_output,
+ apply_conditioning_mask=True,
+ )
+ finally:
+ self._clear_reference_buffers(normal_output)
+
+ return self.teacher_model.loss(
+ normal_batch,
+ model_output,
+ apply_conditioning_mask=True,
+ )
+
+ def _normal_h3_batch(self, prepared_batch: Dict[str, Any]) -> Dict[str, Any]:
+ normal_batch = dict(prepared_batch)
+ for key in self._inner_timestep_conditioning_keys():
+ normal_batch.pop(key, None)
+
+ video_target = self._normal_video_target(prepared_batch)
+ normal_batch["target"] = video_target
+ normal_batch["flow_target"] = video_target
+
+ audio_target = self._normal_audio_target(prepared_batch)
+ if audio_target is None:
+ normal_batch.pop("audio_target", None)
+ else:
+ normal_batch["audio_target"] = audio_target
+ return normal_batch
+
+ def _normal_video_target(self, prepared_batch: Dict[str, Any]) -> torch.Tensor:
+ get_target = getattr(self.teacher_model, "get_flow_matching_target", None)
+ if not callable(get_target):
+ raise ValueError("H3 drift requires MiniMax-H3 get_flow_matching_target() to compute normal SFT loss.")
+
+ target = get_target(prepared_batch, prefer_explicit_target=False)
+ if not torch.is_tensor(target):
+ raise ValueError(f"H3 drift normal video target must be a tensor, got {type(target)}.")
+ return target.detach()
+
+ @staticmethod
+ def _normal_audio_target(prepared_batch: Dict[str, Any]) -> Optional[torch.Tensor]:
+ audio_target = prepared_batch.get("audio_target")
+ if audio_target is not None:
+ if not torch.is_tensor(audio_target):
+ raise ValueError(f"H3 drift audio_target must be a tensor, got {type(audio_target)}.")
+ return audio_target.detach()
+
+ audio_latents = prepared_batch.get("audio_latents")
+ audio_noise = prepared_batch.get("audio_noise")
+ if audio_latents is None and audio_noise is None:
+ return None
+ if not torch.is_tensor(audio_latents) or not torch.is_tensor(audio_noise):
+ raise ValueError("H3 drift audio SFT loss requires tensor audio_latents and audio_noise.")
+ return (audio_latents - audio_noise).detach()
+
+ def _has_inner_timestep_conditioning(self, prepared_batch: Dict[str, Any]) -> bool:
+ return any(prepared_batch.get(key) is not None for key in self._inner_timestep_conditioning_keys())
+
+ def _inner_timestep_conditioning_keys(self) -> tuple[str, ...]:
+ keys = list(self._FLOWMAP_BATCH_KEYS)
+ flowmap_key = getattr(self.teacher_model, "FLOWMAP_R_TIMESTEP_BATCH_KEY", None)
+ if isinstance(flowmap_key, str) and flowmap_key not in keys:
+ keys.append(flowmap_key)
+ return tuple(keys)
+
+ @staticmethod
+ def _inner_logs(logs: Dict[str, Any]) -> Dict[str, float]:
+ inner_logs: Dict[str, float] = {}
+ for key, value in logs.items():
+ log_key = "h3_drift_inner_total" if key == "total" else key
+ if torch.is_tensor(value):
+ value = float(value.detach())
+ inner_logs[log_key] = float(value)
+ return inner_logs
+
+ def generator_loss_step(
+ self,
+ prepared_batch: Dict[str, Any],
+ model_output: Dict[str, Any],
+ current_loss: torch.Tensor,
+ ):
+ if self.inner_distiller is None:
+ return current_loss, {}
+ loss, logs = self.inner_distiller.generator_loss_step(prepared_batch, model_output, current_loss)
+ return loss, self._inner_logs(logs)
+
+ def discriminator_step(self, prepared_batch: Dict[str, Any], **kwargs):
+ if self.inner_distiller is not None:
+ return self.inner_distiller.discriminator_step(prepared_batch=prepared_batch, **kwargs)
+ return None
+
+ def on_load_checkpoint(self, ckpt_dir: str):
+ if self.inner_distiller is not None:
+ return self.inner_distiller.on_load_checkpoint(ckpt_dir)
+ return None
+
+ def on_save_checkpoint(self, step: int, ckpt_dir: str):
+ if self.inner_distiller is not None:
+ return self.inner_distiller.on_save_checkpoint(step, ckpt_dir)
+ return None
+
+ def on_epoch_end(self, epoch: int):
+ if self.inner_distiller is not None:
+ return self.inner_distiller.on_epoch_end(epoch)
+ return None
+
+ @staticmethod
+ def _prediction_from_output(output: Dict[str, Any]) -> _H3Prediction:
+ video = output.get("model_prediction")
+ audio = output.get("audio_prediction")
+ if video is not None and not torch.is_tensor(video):
+ raise ValueError(f"H3 drift video prediction must be a tensor or None, got {type(video)}.")
+ if audio is not None and not torch.is_tensor(audio):
+ raise ValueError(f"H3 drift audio prediction must be a tensor or None, got {type(audio)}.")
+ return _H3Prediction(video=video, audio=audio)
+
+ @staticmethod
+ def _clear_reference_buffers(output: Dict[str, Any]) -> None:
+ hidden_states_buffer = output.get("hidden_states_buffer")
+ if isinstance(hidden_states_buffer, dict):
+ hidden_states_buffer.clear()
+
+ @staticmethod
+ def _broadcast_mask(mask: Optional[torch.Tensor], target: torch.Tensor) -> torch.Tensor:
+ if mask is None:
+ return torch.ones_like(target, dtype=torch.bool)
+ if not torch.is_tensor(mask):
+ raise ValueError(f"H3 drift loss mask must be a tensor, got {type(mask)}.")
+ mask = mask.to(device=target.device, dtype=torch.bool)
+ if mask.shape == target.shape:
+ return mask
+ if mask.ndim == 1:
+ if mask.shape[0] != target.shape[0]:
+ raise ValueError(f"H3 drift loss mask batch {mask.shape[0]} does not match target batch {target.shape[0]}.")
+ mask = mask.view(mask.shape[0], *([1] * (target.ndim - 1)))
+ elif mask.ndim < target.ndim:
+ if mask.shape[0] != target.shape[0]:
+ raise ValueError(f"H3 drift loss mask batch {mask.shape[0]} does not match target batch {target.shape[0]}.")
+ mask = mask.view(mask.shape[0], *([1] * (target.ndim - mask.ndim)), *mask.shape[1:])
+ try:
+ return mask.expand_as(target)
+ except RuntimeError as exc:
+ raise ValueError(
+ f"H3 drift loss mask shape {tuple(mask.shape)} cannot broadcast to {tuple(target.shape)}."
+ ) from exc
+
+ @staticmethod
+ def _modality_loss(
+ prediction: torch.Tensor,
+ target: torch.Tensor,
+ mask: Optional[torch.Tensor],
+ sample_weight: Optional[torch.Tensor],
+ ) -> tuple[torch.Tensor, torch.Tensor, int]:
+ if prediction.shape != target.shape:
+ raise ValueError(
+ f"H3 drift prediction shape {tuple(prediction.shape)} does not match reference {tuple(target.shape)}."
+ )
+ valid = H3DriftDistiller._broadcast_mask(mask, target)
+ elements = int(valid.sum().item())
+ if elements == 0:
+ zero = prediction.sum() * 0.0
+ return zero, zero, 0
+
+ squared = (prediction.float() - target.detach().float()).square()
+ if sample_weight is not None:
+ if not torch.is_tensor(sample_weight):
+ raise ValueError(f"H3 drift sample_weight must be a tensor, got {type(sample_weight)}.")
+ if sample_weight.shape != (target.shape[0],):
+ raise ValueError("H3 drift sample_weight must contain one value per batch item.")
+ weight = sample_weight.to(device=squared.device, dtype=squared.dtype)
+ squared = squared * weight.view(weight.shape[0], *([1] * (squared.ndim - 1)))
+ total = squared.masked_select(valid).sum()
+ return total / elements, total, elements
+
+ @staticmethod
+ def _joint_prediction_loss(
+ prediction: _H3Prediction,
+ target: _H3Prediction,
+ *,
+ video_mask: Optional[torch.Tensor] = None,
+ audio_mask: Optional[torch.Tensor] = None,
+ sample_weight: Optional[torch.Tensor] = None,
+ balance: str = "token",
+ video_weight: float = 1.0,
+ audio_weight: float = 1.0,
+ ) -> _H3JointLoss:
+ if balance not in {"token", "modality"}:
+ raise ValueError(f"Unsupported H3 drift loss balance: {balance}.")
+ if video_weight < 0 or audio_weight < 0 or video_weight + audio_weight <= 0:
+ raise ValueError("H3 drift video/audio loss weights must be non-negative and not both zero.")
+
+ zero_source = prediction.video if prediction.video is not None else prediction.audio
+ if zero_source is None:
+ raise ValueError("H3 drift prediction contains no target modality.")
+ zero = zero_source.sum() * 0.0
+
+ if prediction.video is None or target.video is None:
+ if prediction.video is not None or target.video is not None:
+ raise ValueError("H3 drift video prediction and reference presence differ.")
+ video_mean, video_total, video_elements = zero, zero, 0
+ else:
+ video_mean, video_total, video_elements = H3DriftDistiller._modality_loss(
+ prediction.video, target.video, video_mask, sample_weight
+ )
+
+ if prediction.audio is None or target.audio is None:
+ if prediction.audio is not None or target.audio is not None:
+ raise ValueError("H3 drift audio prediction and reference presence differ.")
+ audio_mean, audio_total, audio_elements = zero, zero, 0
+ else:
+ audio_mean, audio_total, audio_elements = H3DriftDistiller._modality_loss(
+ prediction.audio, target.audio, audio_mask, sample_weight
+ )
+
+ active_video_weight = video_weight if video_elements else 0.0
+ active_audio_weight = audio_weight if audio_elements else 0.0
+ if active_video_weight + active_audio_weight == 0:
+ raise ValueError("H3 drift loss masks exclude every video and audio element.")
+
+ if balance == "modality":
+ loss = (active_video_weight * video_mean + active_audio_weight * audio_mean) / (
+ active_video_weight + active_audio_weight
+ )
+ else:
+ weighted_elements = active_video_weight * video_elements + active_audio_weight * audio_elements
+ loss = (active_video_weight * video_total + active_audio_weight * audio_total) / weighted_elements
+
+ return _H3JointLoss(loss, video_mean, audio_mean, video_elements, audio_elements)
+
+ @staticmethod
+ def _video_mask_for_loss(
+ prepared_batch: Dict[str, Any],
+ prediction: Optional[torch.Tensor],
+ ) -> Optional[torch.Tensor]:
+ if prediction is None or prepared_batch.get("loss_mask_type") not in {"mask", "segmentation"}:
+ return None
+ mask_image = prepared_batch.get("conditioning_pixel_values")
+ if isinstance(mask_image, list):
+ mask_image = mask_image[-1] if mask_image else None
+ if not torch.is_tensor(mask_image):
+ return None
+
+ mask_image = mask_image.to(device=prediction.device, dtype=prediction.dtype)
+ if mask_image.dim() == 3:
+ mask_image = mask_image.unsqueeze(1)
+ if mask_image.dim() == 4:
+ if prepared_batch.get("loss_mask_type") == "segmentation":
+ mask_image = torch.sum(mask_image, dim=1, keepdim=True) / mask_image.shape[1]
+ elif mask_image.shape[1] > 1:
+ mask_image = mask_image[:, 0:1]
+ if prediction.dim() == 5:
+ mask_image = mask_image.unsqueeze(2)
+ elif mask_image.dim() == 5:
+ if prepared_batch.get("loss_mask_type") == "segmentation":
+ mask_image = torch.sum(mask_image, dim=1, keepdim=True) / mask_image.shape[1]
+ elif mask_image.shape[1] > 1:
+ mask_image = mask_image[:, 0:1]
+
+ if mask_image.dim() != prediction.dim():
+ raise ValueError(
+ f"H3 drift video mask rank must match prediction rank. Got {mask_image.dim()} vs {prediction.dim()}."
+ )
+ mask_image = F.interpolate(mask_image, size=prediction.shape[2:], mode="area")
+ mask_image = (mask_image / 2 + 0.5).clamp(0.0, 1.0)
+ if prepared_batch.get("loss_mask_type") == "segmentation":
+ mask_image = (mask_image > 0).to(dtype=prediction.dtype)
+ return mask_image
+
+
+DistillationRegistry.register(
+ "h3_drift",
+ H3DriftDistiller,
+ requires_distillation_cache=False,
+ requirement_notes="MiniMax-H3 LoRA/LyCORIS-only base-prediction drift regularizer.",
+)
diff --git a/simpletuner/helpers/distillation/perflow/distiller.py b/simpletuner/helpers/distillation/perflow/distiller.py
index ea5f111fe..41eab8b3b 100644
--- a/simpletuner/helpers/distillation/perflow/distiller.py
+++ b/simpletuner/helpers/distillation/perflow/distiller.py
@@ -249,15 +249,15 @@ def prepare_batch(self, batch: Dict[str, Any], *_):
noisy_latents = self._distillation_scheduler.add_noise(latents.float(), noise.float(), timesteps).to(
device=device, dtype=dtype
)
- flow_target = noise - latents
-
batch["latents"] = latents
batch["clean_latents"] = latents
batch["noise"] = noise
batch["input_noise"] = noise
batch["timesteps"] = timesteps
batch["noisy_latents"] = noisy_latents
+ flow_target = self._flow_target(batch)
batch["flow_target"] = flow_target
+ batch["target"] = flow_target
batch["distillation_metadata"] = [entry.get("metadata", {}) for entry in cache_entries]
batch["distillation_cache_entries"] = cache_entries
return batch
@@ -284,7 +284,7 @@ def compute_distill_loss(
flow_target = prepared_batch.get("flow_target")
if flow_target is None:
- flow_target = prepared_batch["noise"] - prepared_batch["latents"]
+ flow_target = self._flow_target(prepared_batch)
loss_type = str(self.config.get("loss_type", "l2")).lower()
if loss_type in ("huber", "smooth_l1"):
@@ -303,6 +303,14 @@ def compute_distill_loss(
}
return loss, logs
+ def _flow_target(self, prepared_batch: Dict[str, Any]) -> torch.Tensor:
+ latents = prepared_batch["latents"]
+ get_target = getattr(self.student_model, "get_flow_matching_target", None)
+ if callable(get_target):
+ target = get_target(prepared_batch, prefer_explicit_target=False)
+ return target.to(device=latents.device, dtype=latents.dtype)
+ return prepared_batch["noise"] - latents
+
DistillationRegistry.register(
"perflow",
diff --git a/simpletuner/helpers/distillation/requirements.py b/simpletuner/helpers/distillation/requirements.py
index 77548c56b..97cbef6d0 100644
--- a/simpletuner/helpers/distillation/requirements.py
+++ b/simpletuner/helpers/distillation/requirements.py
@@ -138,6 +138,34 @@ def parse_distiller_requirement_profile(metadata: Mapping[str, Any] | None) -> D
)
+def merge_distiller_requirement_profiles(
+ *profiles: Optional[DistillerRequirementProfile],
+) -> DistillerRequirementProfile:
+ """Merge multiple distiller requirement profiles while preserving declaration order."""
+ requirements: list[DataRequirement] = []
+ notes: list[str] = []
+ is_data_generator = False
+
+ for profile in profiles:
+ if profile is None:
+ continue
+ is_data_generator = is_data_generator or profile.is_data_generator
+ for requirement in profile.requirements:
+ if requirement not in requirements:
+ requirements.append(requirement)
+ for note in profile.notes:
+ if note not in notes:
+ notes.append(note)
+
+ if not requirements and not is_data_generator and not notes:
+ return EMPTY_PROFILE
+ return DistillerRequirementProfile(
+ requirements=tuple(requirements),
+ is_data_generator=is_data_generator,
+ notes=tuple(notes),
+ )
+
+
def _entry_disabled(entry: Any) -> bool:
return isinstance(entry, Mapping) and bool(entry.get("disabled") or entry.get("disable"))
diff --git a/simpletuner/helpers/distillation/self_forcing/wrappers.py b/simpletuner/helpers/distillation/self_forcing/wrappers.py
index 0cfb015f1..b366ddadf 100644
--- a/simpletuner/helpers/distillation/self_forcing/wrappers.py
+++ b/simpletuner/helpers/distillation/self_forcing/wrappers.py
@@ -61,9 +61,11 @@ def forward(
prediction = self.foundation.model_predict(prepared)["model_prediction"]
flow_pred = prediction
+ converter = getattr(self.foundation, "prediction_to_noiseward_flow", None)
+ scheduler_flow = converter(flow_pred) if callable(converter) else flow_pred
latent_shape = flow_pred.shape
- flat_flow, _ = _flatten_video(flow_pred)
+ flat_flow, _ = _flatten_video(scheduler_flow)
flat_xt, _ = _flatten_video(noisy_latents)
flat_timestep = timesteps.reshape(-1)
@@ -81,12 +83,13 @@ def forward(
@dataclass
class ModuleWrapper:
"""
- Wraps a raw Wan transformer module (teacher/fake score) to expose the same API.
+ Wraps a raw transformer module (teacher/fake score) to expose the same API.
"""
module: torch.nn.Module
scheduler: FlowMatchingSchedulerAdapter
weight_dtype: torch.dtype
+ foundation: object | None = None
def forward(
self,
@@ -117,10 +120,14 @@ def forward(
kwargs[f"{key}"] = conditional_dict[key].to(dtype=self.weight_dtype)
outputs = self.module(**kwargs)
- flow_pred = outputs[0]
+ raw_prediction = outputs[0]
+ converter = getattr(self.foundation, "raw_model_prediction_to_model_prediction", None)
+ flow_pred = converter(raw_prediction) if callable(converter) else raw_prediction
+ noiseward_converter = getattr(self.foundation, "prediction_to_noiseward_flow", None)
+ scheduler_flow = noiseward_converter(flow_pred) if callable(noiseward_converter) else flow_pred
latent_shape = flow_pred.shape
- flat_flow, _ = _flatten_video(flow_pred)
+ flat_flow, _ = _flatten_video(scheduler_flow)
flat_xt, _ = _flatten_video(latents)
flat_timestep = timesteps.reshape(-1)
diff --git a/simpletuner/helpers/image_manipulation/training_sample.py b/simpletuner/helpers/image_manipulation/training_sample.py
index 79824431d..4911f8429 100644
--- a/simpletuner/helpers/image_manipulation/training_sample.py
+++ b/simpletuner/helpers/image_manipulation/training_sample.py
@@ -149,6 +149,8 @@ def training_sample_path(self, training_dataset_id: str) -> str:
cond_backend = StateTracker.get_data_backend(self.data_backend_id)
if training_backend is None:
raise ValueError(f"No training dataset registered for backend “{training_dataset_id}”.")
+ if isinstance(self.image_metadata, dict) and self.image_metadata.get("training_sample_path"):
+ return training_backend["data_backend"].get_abs_path(self.image_metadata["training_sample_path"])
training_data_dir = training_backend["config"]["instance_data_dir"]
cond_data_dir = cond_backend["config"]["instance_data_dir"]
if os.path.isabs(self._image_path):
diff --git a/simpletuner/helpers/metadata/backends/base.py b/simpletuner/helpers/metadata/backends/base.py
index f4b6f5fcb..e4d0864ca 100644
--- a/simpletuner/helpers/metadata/backends/base.py
+++ b/simpletuner/helpers/metadata/backends/base.py
@@ -191,7 +191,7 @@ def _apply_max_num_samples_limit(self, file_list: list) -> list:
# Use dataset ID as seed for deterministic selection
rng = random.Random(self.id)
- shuffled = list(file_list)
+ shuffled = sorted(file_list, key=str)
rng.shuffle(shuffled)
limited = shuffled[: self.max_num_samples]
logger.info(f"({self.id}) Applied max_num_samples limit: selected {len(limited)} of {len(file_list)} samples")
@@ -781,7 +781,7 @@ def split_buckets_between_processes(self, gradient_accumulation_steps=1, apply_p
backend_config = StateTracker.get_data_backend_config(self.id) or {}
configured_repeats = int(backend_config.get("repeats") or 0)
user_set_repeats = configured_repeats > 0
- auto_repeat_count = None
+ auto_repeat_counts = {}
# Early validation: check if configuration is mathematically impossible
buckets_that_will_fail = []
@@ -806,22 +806,19 @@ def split_buckets_between_processes(self, gradient_accumulation_steps=1, apply_p
needed_repeats = ceil(effective_batch_size / images) - 1
min_repeats_needed[bucket_info["bucket"]] = needed_repeats
- # The documented repeat setting applies to the whole backend, so
- # use the maximum requirement and apply it consistently to every
- # bucket rather than assigning a different repeat count per bucket.
max_needed_repeats = max(min_repeats_needed.values())
allow_oversubscription = StateTracker.get_args().allow_dataset_oversubscription
# Check if dataset oversubscription is allowed
if allow_oversubscription and not user_set_repeats:
- # Automatically adjust repeats to make training possible
- original_repeats = self.repeats
- auto_repeat_count = max_needed_repeats
+ # Pad only undersized buckets. Applying the rarest bucket's repeat
+ # count backend-wide can multiply an otherwise large dataset.
+ auto_repeat_counts = min_repeats_needed
logger.warning(
- f"(id={self.id}) Dataset oversubscription enabled: automatically increasing repeats from {original_repeats} to {auto_repeat_count}\n"
+ f"(id={self.id}) Dataset oversubscription enabled: automatically padding {len(auto_repeat_counts)} undersized bucket(s)\n"
f" - This allows training with {total_samples} samples across {num_processes} GPUs\n"
f" - Effective batch size: {effective_batch_size}\n"
- f" - Logical repeat factor before per-bucket batch padding: {auto_repeat_count + 1}"
+ f" - Maximum per-bucket repeat factor: {max_needed_repeats + 1}"
)
# Validation passed with adjustment, continue
else:
@@ -873,9 +870,10 @@ def split_buckets_between_processes(self, gradient_accumulation_steps=1, apply_p
if should_shuffle_contents:
# Process-specific RNG seeding must not change the split input across ranks.
shuffle_seed = getattr(StateTracker.get_args(), "seed", None)
- if self.accelerator.is_main_process and shuffle_seed is None:
- shuffle_seed = random.SystemRandom().getrandbits(64)
- shuffle_seed = broadcast_object_from_main(shuffle_seed if self.accelerator.is_main_process else None)
+ if shuffle_seed is None:
+ if self.accelerator.is_main_process:
+ shuffle_seed = random.SystemRandom().getrandbits(64)
+ shuffle_seed = broadcast_object_from_main(shuffle_seed if self.accelerator.is_main_process else None)
for bucket, images in self.aspect_ratio_bucket_indices.items():
if not images:
@@ -883,9 +881,12 @@ def split_buckets_between_processes(self, gradient_accumulation_steps=1, apply_p
continue
if should_shuffle_contents:
logger.debug(f"Shuffling bucket {bucket} contents.")
- images = images.copy()
+ # Cache builders and cache reloads can preserve the same set in different
+ # orders. Canonicalize first so every rank shuffles an identical sequence.
+ images = sorted(images, key=str)
random.Random(f"{shuffle_seed}:{self.id}:{bucket}").shuffle(images)
+ auto_repeat_count = auto_repeat_counts.get(bucket)
if auto_repeat_count is not None:
logical_count = len(images) * (auto_repeat_count + 1)
scheduled_count = ceil(logical_count / effective_batch_size) * effective_batch_size
diff --git a/simpletuner/helpers/metadata/backends/webshart.py b/simpletuner/helpers/metadata/backends/webshart.py
index aed68a5dd..6068a3abd 100644
--- a/simpletuner/helpers/metadata/backends/webshart.py
+++ b/simpletuner/helpers/metadata/backends/webshart.py
@@ -1,8 +1,13 @@
import json
import logging
import os
+import shutil
+import subprocess
+import tempfile
import time
+from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
+from pathlib import Path
from typing import Any, Dict, List, Optional, Union
from tqdm import tqdm
@@ -11,6 +16,7 @@
from simpletuner.helpers.data_backend.webshart import WebshartDataBackend
from simpletuner.helpers.image_manipulation.training_sample import TrainingSample
from simpletuner.helpers.metadata.backends.base import MetadataBackend
+from simpletuner.helpers.training import video_file_extensions
from simpletuner.helpers.training.multi_process import should_log
from simpletuner.helpers.training.state_tracker import StateTracker
@@ -82,8 +88,8 @@ def __init__(
)
if not isinstance(data_backend, WebshartDataBackend):
raise ValueError("WebshartMetadataBackend requires WebshartDataBackend")
- if self.dataset_type not in {DatasetType.IMAGE, DatasetType.CONDITIONING, DatasetType.EVAL}:
- raise ValueError("WebshartMetadataBackend currently supports image-like datasets only.")
+ if self.dataset_type not in {DatasetType.IMAGE, DatasetType.VIDEO, DatasetType.CONDITIONING, DatasetType.EVAL}:
+ raise ValueError("WebshartMetadataBackend supports image, video, conditioning, and eval datasets only.")
self.caption_cache: Dict[str, Union[str, List[str]]] = {}
@@ -202,15 +208,101 @@ def _sample_id_from_entry(self, shard_idx: int, entry: dict) -> str:
raise ValueError("Webshart sample bucket entries must include sample_idx.")
return self.data_backend.sample_id(shard_idx, int(entry["sample_idx"]), str(entry["filename"]))
- def _metadata_for_entry(self, shard_metadata: dict, filename: str, entry: dict) -> dict:
+ @staticmethod
+ def _coerce_positive_number(value: Any, value_type):
+ try:
+ result = value_type(value)
+ except (TypeError, ValueError):
+ return None
+ return result if result > 0 else None
+
+ @staticmethod
+ def _ffprobe_video_path(probe_path: str) -> dict:
+ result = subprocess.run(
+ [
+ "ffprobe",
+ "-v",
+ "error",
+ "-select_streams",
+ "v:0",
+ "-show_entries",
+ "stream=width,height,nb_frames,avg_frame_rate,r_frame_rate,duration",
+ "-show_entries",
+ "format=duration",
+ "-of",
+ "json",
+ probe_path,
+ ],
+ capture_output=True,
+ text=True,
+ check=True,
+ timeout=30,
+ )
+ probe = json.loads(result.stdout) if result.stdout else {}
+ streams = probe.get("streams") or []
+ stream = streams[0] if streams else {}
+ width = WebshartMetadataBackend._coerce_positive_number(stream.get("width"), int)
+ height = WebshartMetadataBackend._coerce_positive_number(stream.get("height"), int)
+ metadata = {"original_size": (width, height)} if width and height else {}
+
+ num_frames = WebshartMetadataBackend._coerce_positive_number(stream.get("nb_frames"), int)
+ if num_frames:
+ metadata["num_frames"] = num_frames
+ duration = WebshartMetadataBackend._coerce_positive_number(
+ stream.get("duration") or (probe.get("format") or {}).get("duration"),
+ float,
+ )
+ if duration:
+ metadata["video_duration"] = duration
+ return metadata
+
+ def _probe_video_metadata(self, sample_path: str, file_metadata: Optional[dict] = None) -> dict:
+ if shutil.which("ffprobe") is None:
+ return {}
+
+ suffix = Path(self.data_backend.parse_sample_id(sample_path).filename).suffix or ".mp4"
+ probe_path = None
+ try:
+ range_reader = getattr(self.data_backend, "read_sample_head_tail", None)
+ if callable(range_reader):
+ try:
+ head, tail, total_size = range_reader(sample_path, file_metadata=file_metadata)
+ with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as handle:
+ probe_path = handle.name
+ handle.truncate(total_size)
+ handle.seek(0)
+ handle.write(head)
+ handle.seek(total_size - len(tail))
+ handle.write(tail)
+ metadata = self._ffprobe_video_path(probe_path)
+ if metadata.get("original_size"):
+ return metadata
+ except Exception as exc:
+ logger.debug("Unable to range-probe Webshart video %s: %s", sample_path, exc)
+ finally:
+ if probe_path:
+ Path(probe_path).unlink(missing_ok=True)
+ probe_path = None
+
+ payload = self.data_backend.read(sample_path)
+ if not payload:
+ return {}
+ with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as handle:
+ handle.write(payload)
+ probe_path = handle.name
+ return self._ffprobe_video_path(probe_path)
+ except Exception as exc:
+ logger.debug("Unable to probe Webshart video %s: %s", sample_path, exc)
+ return {}
+ finally:
+ if probe_path:
+ Path(probe_path).unlink(missing_ok=True)
+
+ def _metadata_for_entry(self, shard_metadata: dict, filename: str, entry: dict, sample_path: str) -> dict:
file_metadata = shard_metadata.get(filename, {}) or {}
width = entry.get("width", file_metadata.get("width"))
height = entry.get("height", file_metadata.get("height"))
- if width is None or height is None:
- return {}
-
metadata = {
- "original_size": (int(width), int(height)),
"webshart": {
"shard_idx": entry.get("shard_idx"),
"sample_idx": entry.get("sample_idx"),
@@ -220,14 +312,51 @@ def _metadata_for_entry(self, shard_metadata: dict, filename: str, entry: dict)
"json_path": file_metadata.get("json_path"),
},
}
+ if width is not None and height is not None:
+ metadata["original_size"] = (int(width), int(height))
if "captions" in file_metadata:
metadata["captions"] = file_metadata["captions"]
- if "json_metadata" in file_metadata:
- metadata["json_metadata"] = file_metadata["json_metadata"]
+ json_metadata = file_metadata.get("json_metadata") or {}
+ if json_metadata:
+ metadata["json_metadata"] = json_metadata
if "json_path" in file_metadata:
metadata["json_path"] = file_metadata["json_path"]
+
+ if self.dataset_type is DatasetType.VIDEO:
+ fps = self._coerce_positive_number(json_metadata.get("fps"), float)
+ num_frames = self._coerce_positive_number(
+ json_metadata.get("frame", json_metadata.get("num_frames")),
+ int,
+ )
+ duration = self._coerce_positive_number(
+ json_metadata.get("seconds", json_metadata.get("duration")),
+ float,
+ )
+ if fps:
+ metadata["fps"] = fps
+ if num_frames:
+ metadata["num_frames"] = num_frames
+ if duration:
+ metadata["video_duration"] = duration
+ if "original_size" not in metadata:
+ probed = self._probe_video_metadata(sample_path, file_metadata=file_metadata)
+ for key, value in probed.items():
+ metadata.setdefault(key, value)
return metadata
+ def _prepare_bucket_entry(
+ self,
+ shard_metadata: dict,
+ entry: dict,
+ sample_path: str,
+ ) -> tuple[dict, Optional[tuple[float, dict]], Optional[Exception]]:
+ try:
+ filename = str(entry["filename"])
+ sample_metadata = self._metadata_for_entry(shard_metadata, filename, entry, sample_path)
+ return sample_metadata, self._prepare_metadata(sample_path, sample_metadata), None
+ except Exception as exc:
+ return {}, None, exc
+
def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional[tuple[float, dict]]:
if not sample_metadata or "original_size" not in sample_metadata:
return None
@@ -249,7 +378,32 @@ def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional
"target_size": prepared_sample.target_size,
}
)
- return round(aspect_ratio, 2), sample_metadata
+ if self.dataset_type is DatasetType.VIDEO and self.bucket_strategy == "resolution_frames":
+ target_width, target_height = prepared_sample.target_size
+ bucket_key, rounded_frames = self._compute_video_bucket(
+ target_width,
+ target_height,
+ sample_metadata["num_frames"],
+ )
+ sample_metadata["bucket_frames"] = rounded_frames
+ else:
+ bucket_key = round(aspect_ratio, 2)
+ return bucket_key, sample_metadata
+
+ def _entries_for_shard(self, shard_idx: int) -> list[dict]:
+ if self.dataset_type is not DatasetType.VIDEO:
+ shard_bucket_results = self.data_backend.list_shard_sample_aspect_buckets([shard_idx])
+ if not shard_bucket_results:
+ return []
+ shard_bucket_data = shard_bucket_results[0]
+ return [entry for entries in shard_bucket_data.get("buckets", {}).values() for entry in entries]
+
+ entries = []
+ for sample_idx, filename in enumerate(self.data_backend.dataset.list_samples_in_shard(shard_idx)):
+ if Path(filename).suffix.lower().strip(".") not in video_file_extensions:
+ continue
+ entries.append({"sample_idx": sample_idx, "filename": filename})
+ return entries
def compute_aspect_ratio_bucket_indices(self, ignore_existing_cache: bool = False, progress_callback=None):
logger.info("Building aspect ratio buckets from webshart metadata...")
@@ -282,42 +436,71 @@ def compute_aspect_ratio_bucket_indices(self, ignore_existing_cache: bool = Fals
metadata_updates: Dict[str, dict] = {}
shard_metadata_cache: Dict[int, dict] = {}
shard_indices = self._all_shard_indices()
+ worker_count = max(1, int(getattr(self.data_backend, "parallel_downloads", 1)))
+ executor = (
+ ThreadPoolExecutor(max_workers=worker_count)
+ if self.dataset_type is DatasetType.VIDEO and worker_count > 1
+ else None
+ )
- for shard_idx in tqdm(
- shard_indices,
- desc="Processing webshart metadata",
- total=len(shard_indices),
- leave=False,
- ncols=100,
- ):
- if (
- self.max_num_samples is not None
- and statistics["total_processed"] + len(existing_files) >= self.max_num_samples
+ try:
+ for shard_idx in tqdm(
+ shard_indices,
+ desc="Processing webshart metadata",
+ total=len(shard_indices),
+ leave=False,
+ ncols=100,
):
- break
+ if (
+ self.max_num_samples is not None
+ and statistics["total_processed"] + len(existing_files) >= self.max_num_samples
+ ):
+ break
- shard_bucket_results = self.data_backend.list_shard_sample_aspect_buckets([shard_idx])
- if not shard_bucket_results:
- continue
- shard_bucket_data = shard_bucket_results[0]
- shard_metadata = shard_metadata_cache.setdefault(shard_idx, self.data_backend.get_shard_metadata(shard_idx))
- for entries in shard_bucket_data.get("buckets", {}).values():
- for entry in entries:
+ shard_entries = self._entries_for_shard(shard_idx)
+ if not shard_entries:
+ continue
+ shard_metadata = shard_metadata_cache.setdefault(shard_idx, self.data_backend.get_shard_metadata(shard_idx))
+ candidates = []
+ for entry in shard_entries:
+ entry = {**entry, "shard_idx": shard_idx}
+ sample_path = self._sample_id_from_entry(shard_idx, entry)
+ if sample_path in existing_files:
+ continue
+ candidates.append((entry, sample_path))
+
+ chunk_size = max(1, worker_count * 2)
+ candidate_idx = 0
+ while candidate_idx < len(candidates):
if (
self.max_num_samples is not None
and statistics["total_processed"] + len(existing_files) >= self.max_num_samples
):
break
- processed_entries += 1
- try:
- filename = str(entry["filename"])
- entry = {**entry, "shard_idx": shard_idx}
- sample_path = self._sample_id_from_entry(shard_idx, entry)
- if sample_path in existing_files:
+ remaining = (
+ self.max_num_samples - statistics["total_processed"] - len(existing_files)
+ if self.max_num_samples is not None
+ else chunk_size
+ )
+ current_chunk_size = min(chunk_size, remaining)
+ chunk = candidates[candidate_idx : candidate_idx + current_chunk_size]
+ candidate_idx += len(chunk)
+ if executor is None:
+ results = [
+ self._prepare_bucket_entry(shard_metadata, entry, sample_path) for entry, sample_path in chunk
+ ]
+ else:
+ results = executor.map(
+ lambda item: self._prepare_bucket_entry(shard_metadata, item[0], item[1]),
+ chunk,
+ )
+
+ for (entry, sample_path), (sample_metadata, prepared, error) in zip(chunk, results):
+ processed_entries += 1
+ if error is not None:
+ logger.error("Error processing webshart bucket entry %s: %s", entry, error)
+ statistics["skipped"]["error"] += 1
continue
-
- sample_metadata = self._metadata_for_entry(shard_metadata, filename, entry)
- prepared = self._prepare_metadata(sample_path, sample_metadata)
if prepared is None:
if sample_metadata:
statistics["skipped"]["too_small"] += 1
@@ -330,9 +513,6 @@ def compute_aspect_ratio_bucket_indices(self, ignore_existing_cache: bool = Fals
if sample_metadata.get("captions"):
self.caption_cache[sample_path] = sample_metadata["captions"]
statistics["total_processed"] += 1
- except Exception as exc:
- logger.error("Error processing webshart bucket entry %s: %s", entry, exc)
- statistics["skipped"]["error"] += 1
current_time = time.time()
if (current_time - last_save_time) >= self.metadata_update_interval:
@@ -346,13 +526,11 @@ def compute_aspect_ratio_bucket_indices(self, ignore_existing_cache: bool = Fals
self.save_image_metadata()
self._save_caption_cache()
last_save_time = current_time
- if (
- self.max_num_samples is not None
- and statistics["total_processed"] + len(existing_files) >= self.max_num_samples
- ):
- break
- if progress_callback is not None:
- progress_callback(shard_idx + 1, len(shard_indices))
+ if progress_callback is not None:
+ progress_callback(shard_idx + 1, len(shard_indices))
+ finally:
+ if executor is not None:
+ executor.shutdown(wait=True)
for key, value in aspect_ratio_bucket_updates.items():
self.aspect_ratio_bucket_indices.setdefault(key, []).extend(value)
diff --git a/simpletuner/helpers/metadata/utils/duplicator.py b/simpletuner/helpers/metadata/utils/duplicator.py
index b11816f25..cb169a813 100644
--- a/simpletuner/helpers/metadata/utils/duplicator.py
+++ b/simpletuner/helpers/metadata/utils/duplicator.py
@@ -12,6 +12,23 @@
class DatasetDuplicator:
+ @staticmethod
+ def _translate_conditioning_path(path: str, source_dir: str, target_dir: str, conditioning_data_type: str | None) -> str:
+ source_dir_abs = os.path.abspath(source_dir)
+ target_dir_abs = os.path.abspath(target_dir)
+ if os.path.isabs(path):
+ path_abs = os.path.abspath(path)
+ if os.path.commonpath([path_abs, source_dir_abs]) == source_dir_abs:
+ rel_path = os.path.relpath(path_abs, source_dir_abs)
+ new_path = os.path.join(target_dir_abs, rel_path)
+ else:
+ new_path = os.path.join(target_dir_abs, os.path.basename(path_abs))
+ else:
+ new_path = os.path.join(target_dir, os.path.basename(path))
+ if conditioning_data_type == "i2v_first_frame":
+ new_path = os.path.splitext(new_path)[0] + ".png"
+ return new_path
+
@staticmethod
def copy_metadata(source_backend, target_backend):
"""Copy metadata from source backend to target backend with path updates."""
@@ -27,29 +44,27 @@ def copy_metadata(source_backend, target_backend):
# Get the instance directories for path translation
source_dir = source_backend.get("instance_data_dir", "")
target_dir = target_backend.get("instance_data_dir", "")
+ target_config = target_backend.get("config", {}) or {}
+ conditioning_config = target_config.get("conditioning_config") or target_backend.get("conditioning_config") or {}
+ conditioning_data_type = conditioning_config.get("type")
# Check if we need to update paths (for conditioning datasets)
needs_path_update = source_dir != target_dir and target_backend.get("dataset_type") == "conditioning"
if needs_path_update:
logger.info(f"Copying metadata with path translation: '{source_dir}' -> '{target_dir}'")
- source_dir_abs = os.path.abspath(source_dir)
- target_dir_abs = os.path.abspath(target_dir)
# Copy and update bucket indices
target_meta.aspect_ratio_bucket_indices = {}
for bucket, paths in source_meta.aspect_ratio_bucket_indices.items():
updated_paths = []
for path in paths:
- # Update the path to point to the target directory
- # Handle both absolute and relative paths
- if os.path.isabs(path):
- # For absolute paths, replace the directory
- rel_path = os.path.relpath(path, source_dir_abs)
- new_path = os.path.join(target_dir_abs, rel_path)
- else:
- # For relative paths, just prepend the new directory
- new_path = os.path.join(target_dir, os.path.basename(path))
+ new_path = DatasetDuplicator._translate_conditioning_path(
+ path,
+ source_dir,
+ target_dir,
+ conditioning_data_type,
+ )
updated_paths.append(new_path)
target_meta.aspect_ratio_bucket_indices[bucket] = updated_paths
@@ -57,13 +72,17 @@ def copy_metadata(source_backend, target_backend):
if hasattr(source_meta, "image_metadata") and source_meta.image_metadata:
target_meta.image_metadata = {}
for path, metadata in source_meta.image_metadata.items():
- # Update paths in image metadata too
- if os.path.isabs(path):
- rel_path = os.path.relpath(path, source_dir_abs)
- new_path = os.path.join(target_dir_abs, rel_path)
- else:
- new_path = os.path.join(target_dir, os.path.basename(path))
- target_meta.image_metadata[new_path] = metadata
+ new_path = DatasetDuplicator._translate_conditioning_path(
+ path,
+ source_dir,
+ target_dir,
+ conditioning_data_type,
+ )
+ copied_metadata = dict(metadata)
+ if conditioning_data_type == "i2v_first_frame":
+ copied_metadata["training_sample_path"] = path
+ copied_metadata["image_path"] = new_path
+ target_meta.image_metadata[new_path] = copied_metadata
logger.debug(f"Copied {len(target_meta.image_metadata)} image_metadata entries")
else:
logger.debug("No image_metadata to copy from source")
@@ -79,7 +98,6 @@ def copy_metadata(source_backend, target_backend):
target_meta.set_metadata(metadata_backend=source_meta, update_json=False)
source_config = source_backend.get("config", {}) or {}
- target_config = target_backend.get("config", {}) or {}
conditioning_type = target_config.get("conditioning_type") or target_backend.get("conditioning_type")
propagated_fields = [
@@ -119,6 +137,9 @@ def copy_metadata(source_backend, target_backend):
target_meta.config = target_config
+ if conditioning_data_type == "i2v_first_frame" and hasattr(target_meta, "save_cache"):
+ target_meta.save_cache()
+
# Bucket indices may be rank-local here; do not overwrite the canonical target cache.
target_meta.set_readonly()
if hasattr(target_meta, "save_image_metadata"):
@@ -209,6 +230,8 @@ def _create_single_conditioning_config(source_cfg, cond_cfg, source_id, global_c
target_cfg.pop("audio", None)
target_cfg.pop("s2v_datasets", None)
target_cfg.pop("_s2v_audio_autoinjected", None)
+ if conditioning_data_type == "i2v_first_frame":
+ target_cfg.pop("video", None)
# Set core fields
target_cfg["auto_generated"] = True
diff --git a/simpletuner/helpers/models/boogu_image/model.py b/simpletuner/helpers/models/boogu_image/model.py
index 473489878..9591ed764 100644
--- a/simpletuner/helpers/models/boogu_image/model.py
+++ b/simpletuner/helpers/models/boogu_image/model.py
@@ -173,10 +173,13 @@ def sample_flow_sigmas(self, batch: dict, state: dict) -> tuple[torch.Tensor, to
boogu_timesteps = 1.0 - noise_sigmas
return noise_sigmas, boogu_timesteps
+ def flow_matching_target_direction(self) -> float:
+ return -1.0
+
def get_prediction_target(self, prepared_batch: dict):
if prepared_batch.get("target") is not None:
return prepared_batch["target"]
- return prepared_batch["latents"] - prepared_batch["noise"]
+ return self.get_flow_matching_target(prepared_batch, prefer_explicit_target=False)
def text_embed_cache_key(self):
if self._is_edit_flavour():
diff --git a/simpletuner/helpers/models/chroma/pipeline.py b/simpletuner/helpers/models/chroma/pipeline.py
index e4c5404f1..d28bfd6ca 100644
--- a/simpletuner/helpers/models/chroma/pipeline.py
+++ b/simpletuner/helpers/models/chroma/pipeline.py
@@ -29,7 +29,6 @@
convert_unet_state_dict_to_peft,
deprecate,
get_adapter_name,
- get_peft_kwargs,
is_peft_available,
is_peft_version,
is_torch_version,
@@ -45,6 +44,7 @@
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, T5EncoderModel, T5TokenizerFast
from simpletuner.helpers.models.chroma.transformer import ChromaTransformer2DModel
+from simpletuner.helpers.training.lora_format import get_peft_kwargs
from simpletuner.helpers.training.lycoris import apply_tlora_inference_mask, clear_tlora_mask
from simpletuner.helpers.utils.offloading import restore_offload_state, unpack_offload_state
diff --git a/simpletuner/helpers/models/common.py b/simpletuner/helpers/models/common.py
index f887cd41e..b6a118539 100644
--- a/simpletuner/helpers/models/common.py
+++ b/simpletuner/helpers/models/common.py
@@ -53,11 +53,14 @@
from simpletuner.helpers.training.layersync import LayerSyncRegularizer
from simpletuner.helpers.training.lora_format import (
PEFTLoRAFormat,
+ collect_lora_alphas,
convert_comfyui_to_diffusers,
convert_diffusers_to_comfyui,
convert_diffusers_to_comfyui_sd_lora,
detect_state_dict_format,
normalize_lora_format,
+ peft_lora_config_kwargs_from_state_dict,
+ synthesize_missing_lora_alphas_from_ranks,
)
from simpletuner.helpers.training.min_snr_gamma import compute_snr
from simpletuner.helpers.training.multi_process import _get_rank
@@ -116,6 +119,7 @@ def _is_hf_repo_id(path: str) -> bool:
"zlab_i1",
"ideogram",
"krea2",
+ "minimaxh3",
]
upstream_config_sources = {
"sdxl": "stabilityai/stable-diffusion-xl-base-1.0",
@@ -133,6 +137,7 @@ def _is_hf_repo_id(path: str) -> bool:
"hunyuanvideo": "tencent/HunyuanVideo-1.5",
"ideogram": "ideogram-ai/ideogram-4-fp8",
"krea2": "krea/Krea-2-Raw",
+ "minimaxh3": "MiniMaxAI/MiniMax-H3",
}
@@ -951,9 +956,35 @@ def get_lora_target_layers(self):
combined.append(target)
return combined
+ def _prepare_init_lora_state_dict(self, state_dict: dict) -> dict:
+ return state_dict
+
+ def _load_init_lora_state_dict(self) -> dict | None:
+ init_lora_path = getattr(self.config, "init_lora", None)
+ if not init_lora_path or not isinstance(init_lora_path, str) or not os.path.isfile(init_lora_path):
+ return None
+
+ import safetensors.torch
+
+ try:
+ state_dict = safetensors.torch.load_file(init_lora_path)
+ except Exception as exc:
+ raise ValueError(f"Unable to inspect init_lora checkpoint `{init_lora_path}` for LoRA rank metadata.") from exc
+ return self._prepare_init_lora_state_dict(state_dict)
+
+ def _init_lora_config_kwargs(self, state_dict: dict | None = None) -> dict:
+ state_dict = self._load_init_lora_state_dict() if state_dict is None else state_dict
+ if not state_dict:
+ return {}
+ key_to_replace = (
+ self.CONTROLNET_LORA_STATE_DICT_PREFIX if getattr(self.config, "controlnet", False) else self.MODEL_TYPE.value
+ )
+ return peft_lora_config_kwargs_from_state_dict(state_dict, prefix_to_strip=f"{key_to_replace}.")
+
def add_lora_adapter(self):
from peft import LoraConfig
+ init_lora_state_dict = self._load_init_lora_state_dict()
target_modules = self.get_lora_target_layers()
save_modules = self.get_lora_save_layers()
addkeys, misskeys = [], []
@@ -989,9 +1020,15 @@ def add_lora_adapter(self):
logger.info("Enabling SingLoRA for LoRA training.")
setup_singlora()
+ lora_config_kwargs.update(self._init_lora_config_kwargs(init_lora_state_dict))
+ lora_rank = lora_config_kwargs.pop("r", self.config.lora_rank)
+ lora_alpha = lora_config_kwargs.pop(
+ "lora_alpha",
+ self.config.lora_alpha if self.config.lora_alpha is not None else lora_rank,
+ )
self.lora_config = lora_config_cls(
- r=self.config.lora_rank,
- lora_alpha=(self.config.lora_alpha if self.config.lora_alpha is not None else self.config.lora_rank),
+ r=lora_rank,
+ lora_alpha=lora_alpha,
lora_dropout=self.config.lora_dropout,
init_lora_weights=self.config.lora_initialisation_style,
target_modules=target_modules,
@@ -1020,6 +1057,7 @@ def add_lora_adapter(self):
{self.MODEL_TYPE.value: (self.controlnet if getattr(self.config, "controlnet", False) else self.model)},
self.config.init_lora,
use_dora=use_dora,
+ state_dict=init_lora_state_dict,
)
return addkeys, misskeys
@@ -1526,6 +1564,10 @@ def _prepare_crepa_self_flow_batch(self, batch: dict, state: dict) -> dict:
def uses_audio_latents(self) -> bool:
return False
+ def uses_audio_latents_for_data_backend(self, data_backend_id: Optional[str] = None) -> bool:
+ del data_backend_id
+ return self.uses_audio_latents()
+
def uses_audio_tokens(self) -> bool:
"""
Override to True for autoregressive audio models that consume discrete token sequences
@@ -1865,6 +1907,15 @@ def _normalise_alpha_map(alpha_dict: dict) -> dict:
from diffusers.utils import convert_unet_state_dict_to_peft
denoiser_sd = convert_unet_state_dict_to_peft(denoiser_sd)
+ if not alpha_map:
+ explicit_alphas = collect_lora_alphas(denoiser_sd)
+ if explicit_alphas:
+ alpha_map.update(
+ _normalise_alpha_map({f"{module_key}.alpha": alpha for module_key, alpha in explicit_alphas.items()})
+ )
+ if not alpha_map:
+ alpha_map.update(_normalise_alpha_map(synthesize_missing_lora_alphas_from_ranks(denoiser_sd)))
+ denoiser_sd = {key: value for key, value in denoiser_sd.items() if not key.endswith((".alpha", ".lora_alpha"))}
from peft.utils import set_peft_model_state_dict
@@ -4254,6 +4305,63 @@ def setup_diff2flow_bridge(self):
self.diff2flow_bridge = DiffusionToFlowBridge(alphas_cumprod=alphas_cumprod)
self.diff2flow_bridge.to(device=self.accelerator.device, dtype=self.config.weight_dtype)
+ def flow_matching_target_direction(self) -> float:
+ """
+ Return +1 for the common noise-ward velocity convention (`noise - latents`),
+ or -1 for models trained on the inverse data-ward convention.
+ """
+ return 1.0
+
+ def raw_model_prediction_to_model_prediction(self, raw_prediction: torch.Tensor) -> torch.Tensor:
+ """
+ Convert the denoiser's raw output into SimpleTuner's public `model_prediction`
+ convention. Most models already emit the public convention.
+ """
+ return raw_prediction
+
+ def noiseward_flow_to_prediction(self, flow: torch.Tensor) -> torch.Tensor:
+ direction = float(self.flow_matching_target_direction())
+ if direction == 1.0:
+ return flow
+ if direction == -1.0:
+ return -flow
+ return flow * direction
+
+ def prediction_to_noiseward_flow(self, prediction: torch.Tensor) -> torch.Tensor:
+ direction = float(self.flow_matching_target_direction())
+ if direction == 1.0:
+ return prediction
+ if direction == -1.0:
+ return -prediction
+ if direction == 0.0:
+ raise ValueError("Flow-matching target direction may not be zero.")
+ return prediction / direction
+
+ def flow_matching_target(self, latents: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+ return self.noiseward_flow_to_prediction(noise - latents)
+
+ def get_flow_matching_target(
+ self,
+ prepared_batch: dict,
+ *,
+ latents: Optional[torch.Tensor] = None,
+ noise: Optional[torch.Tensor] = None,
+ prefer_explicit_target: bool = True,
+ ) -> torch.Tensor:
+ if prefer_explicit_target:
+ target = prepared_batch.get("target")
+ if target is not None:
+ return target
+ flow_target = prepared_batch.get("flow_target")
+ if flow_target is not None:
+ return flow_target
+
+ if latents is None:
+ latents = prepared_batch["latents"]
+ if noise is None:
+ noise = prepared_batch["noise"]
+ return self.flow_matching_target(latents, noise)
+
def get_prediction_target(self, prepared_batch: dict):
"""
Returns the target used in the loss function.
@@ -4264,7 +4372,7 @@ def get_prediction_target(self, prepared_batch: dict):
# Parent-student training
target = prepared_batch["target"]
elif self.PREDICTION_TYPE is PredictionTypes.FLOW_MATCHING:
- target = prepared_batch["noise"] - prepared_batch["latents"]
+ target = self.get_flow_matching_target(prepared_batch, prefer_explicit_target=False)
elif self.PREDICTION_TYPE is PredictionTypes.EPSILON:
target = prepared_batch["noise"]
elif self.PREDICTION_TYPE is PredictionTypes.V_PREDICTION:
@@ -4610,6 +4718,19 @@ def sample_flow_sigmas(self, batch: dict, state: dict) -> tuple[torch.Tensor, to
timesteps = sigmas * 1000.0
return sigmas, timesteps
+ def flow_matching_timesteps_from_sigmas(
+ self,
+ sigmas: torch.Tensor,
+ *,
+ reference_timesteps: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Convert noise-ward flow sigmas to the timestep convention consumed by this model."""
+ if reference_timesteps is not None and torch.max(reference_timesteps.detach().float()) <= 1.0:
+ return sigmas
+ scheduler_config = getattr(getattr(self, "noise_schedule", None), "config", None)
+ num_train_timesteps = float(getattr(scheduler_config, "num_train_timesteps", 1000) or 1000)
+ return sigmas * num_train_timesteps
+
def _validate_twinflow_config(self) -> None:
"""
Validate TwinFlow configuration and record common flags.
@@ -4885,8 +5006,8 @@ def _twinflow_generate_fake_samples(
use_grad=False,
)
- # Integrate one step: x_fake = z - F_fake (from t=1 to t=0)
- x_fake = z - F_fake
+ # Integrate one step from t=1 to t=0 using the canonical noise-ward velocity.
+ x_fake = z - self.prediction_to_noiseward_flow(F_fake)
return x_fake.detach(), z
def _twinflow_compute_adversarial_loss(
@@ -4913,8 +5034,8 @@ def _twinflow_compute_adversarial_loss(
# Construct fake trajectory interpolation
x_t_fake = t_b * z + (1 - t_b) * x_fake
- # Target velocity for fake trajectory
- target_fake = z - x_fake
+ # Target velocity for fake trajectory in this model's prediction convention.
+ target_fake = self.noiseward_flow_to_prediction(z - x_fake)
# Forward with NEGATIVE time (sign embedding handles distinction)
neg_t = -t
@@ -5176,15 +5297,15 @@ def _twinflow_forward(
return pred
- @staticmethod
- def _twinflow_reconstruct_states(x_t: torch.Tensor, sigma: torch.Tensor, flow_pred: torch.Tensor):
+ def _twinflow_reconstruct_states(self, x_t: torch.Tensor, sigma: torch.Tensor, flow_pred: torch.Tensor):
"""
Recover x_hat and z_hat from the predicted flow under linear interpolation x_t = t*z + (1-t)*x.
"""
sigma_b = ModelFoundation._twinflow_match_time_shape(sigma, x_t)
gamma = 1 - sigma_b
- x_hat = x_t - sigma_b * flow_pred
- z_hat = x_t + gamma * flow_pred
+ noiseward_flow = self.prediction_to_noiseward_flow(flow_pred)
+ x_hat = x_t - sigma_b * noiseward_flow
+ z_hat = x_t + gamma * noiseward_flow
return x_hat, z_hat
def _twinflow_rcgm_target(
@@ -5330,7 +5451,9 @@ def prepare_batch(self, batch: dict, state: dict) -> dict:
batch["noise"] = noise
if getattr(self.config, "diff2flow_enabled", False):
- batch["flow_target"] = (batch["noise"] - batch["latents"]).to(**target_device_kwargs)
+ batch["flow_target"] = self.get_flow_matching_target(batch, prefer_explicit_target=False).to(
+ **target_device_kwargs
+ )
# Possibly add input perturbation to input noise only
if self.config.input_perturbation != 0 and (
@@ -5460,6 +5583,19 @@ def requires_validation_i2v_samples(self) -> bool:
"""
return False
+ def uses_validation_negative_prompt(self) -> bool:
+ """
+ Whether validation should encode and pass a negative prompt branch.
+ """
+ return self.VALIDATION_USES_NEGATIVE_PROMPT
+
+ def validation_negative_prompt_requires_prompt_context(self) -> bool:
+ """
+ Whether validation negative prompts must be encoded per sample with the same
+ image/reference context as the positive prompt.
+ """
+ return False
+
def should_precompute_validation_negative_prompt(self) -> bool:
"""
Whether to pre-encode negative prompts during validation setup.
@@ -5881,7 +6017,9 @@ def _compute_twinflow_losses(self, prepared_batch: dict, base_pred: torch.Tensor
rcgm_latents = (1 - sigmas) * latents + sigmas * noise
prepared_batch["twinflow_tt"] = tt
- target = (noise - latents).to(device=self.accelerator.device, dtype=self.config.weight_dtype)
+ target = self.get_flow_matching_target(prepared_batch, prefer_explicit_target=False).to(
+ device=self.accelerator.device, dtype=self.config.weight_dtype
+ )
if self._twinflow_diffusion_bridge:
if self.diff2flow_bridge is None:
diff --git a/simpletuner/helpers/models/field_registry/minimaxh3.py b/simpletuner/helpers/models/field_registry/minimaxh3.py
new file mode 100644
index 000000000..dab6d6467
--- /dev/null
+++ b/simpletuner/helpers/models/field_registry/minimaxh3.py
@@ -0,0 +1,53 @@
+from simpletuner.simpletuner_sdk.server.services.field_registry.types import ConfigField, FieldType, ImportanceLevel
+
+
+def register_fields(registry) -> None:
+ registry._add_field(
+ ConfigField(
+ name="minimax_h3_reference_mode",
+ arg_name="--minimax_h3_reference_mode",
+ ui_label="MiniMax H3 Reference Mode",
+ field_type=FieldType.SELECT,
+ tab="model",
+ section="model_specific",
+ model_specific=["minimaxh3"],
+ default_value="vanilla",
+ choices=[
+ {"value": "vanilla", "label": "Vanilla"},
+ {"value": "cached_kv", "label": "Cached KV"},
+ ],
+ help_text="Choose MiniMax-H3 reference handling. cached_kv is an experimental no-grad inference mode that treats text plus conditioning rows as static memory and reuses their K/V projections.",
+ tooltip="Vanilla preserves the stock full packed self-attention path. Cached KV is experimental and intended for reference-conditioning inference comparisons.",
+ importance=ImportanceLevel.EXPERIMENTAL,
+ order=37,
+ documentation="OPTIONS.md#--minimax_h3_reference_mode",
+ )
+ )
+ registry._add_field(
+ ConfigField(
+ name="minimax_h3_target_mode",
+ arg_name="--minimax_h3_target_mode",
+ ui_label="MiniMax H3 Target Mode",
+ field_type=FieldType.SELECT,
+ tab="model",
+ section="model_specific",
+ model_specific=["minimaxh3"],
+ default_value="auto",
+ choices=[
+ {"value": "auto", "label": "Auto"},
+ {"value": "video", "label": "Video only"},
+ {"value": "av", "label": "Audio + video"},
+ ],
+ help_text=(
+ "Controls whether MiniMax-H3 includes target audio rows. auto resolves to video-only unless the "
+ "global config or data backend sets minimax_h3_target_mode/h3_target_mode to av."
+ ),
+ tooltip=(
+ "Use video to ignore auto-split or explicit audio backends and save VRAM; use av for joint "
+ "audio-video training or sampling."
+ ),
+ importance=ImportanceLevel.ADVANCED,
+ order=38,
+ documentation="OPTIONS.md#--minimax_h3_target_mode",
+ )
+ )
diff --git a/simpletuner/helpers/models/flux/pipeline.py b/simpletuner/helpers/models/flux/pipeline.py
index 3a5388eff..8d20441dc 100644
--- a/simpletuner/helpers/models/flux/pipeline.py
+++ b/simpletuner/helpers/models/flux/pipeline.py
@@ -42,7 +42,6 @@
convert_state_dict_to_peft,
convert_unet_state_dict_to_peft,
get_adapter_name,
- get_peft_kwargs,
is_peft_available,
is_peft_version,
is_torch_version,
@@ -70,6 +69,7 @@
PEFTLoRAFormat,
convert_comfyui_to_diffusers,
detect_state_dict_format,
+ get_peft_kwargs,
normalize_lora_format,
)
from simpletuner.helpers.training.lycoris import apply_tlora_inference_mask, clear_tlora_mask
diff --git a/simpletuner/helpers/models/flux/pipeline_controlnet.py b/simpletuner/helpers/models/flux/pipeline_controlnet.py
index 0a2dbdef1..0efcf55bc 100644
--- a/simpletuner/helpers/models/flux/pipeline_controlnet.py
+++ b/simpletuner/helpers/models/flux/pipeline_controlnet.py
@@ -39,7 +39,6 @@
convert_state_dict_to_peft,
convert_unet_state_dict_to_peft,
get_adapter_name,
- get_peft_kwargs,
is_peft_available,
is_peft_version,
is_torch_version,
diff --git a/simpletuner/helpers/models/hidream/model.py b/simpletuner/helpers/models/hidream/model.py
index 17b583f1c..46a9b3d65 100644
--- a/simpletuner/helpers/models/hidream/model.py
+++ b/simpletuner/helpers/models/hidream/model.py
@@ -118,6 +118,9 @@ class HiDream(ImageModelFoundation):
def supports_crepa_self_flow(self) -> bool:
return True
+ def raw_model_prediction_to_model_prediction(self, raw_prediction: torch.Tensor) -> torch.Tensor:
+ return raw_prediction * -1
+
def _prepare_crepa_self_flow_batch(self, batch: dict, state: dict) -> dict:
patch_size = int(max(getattr(self.unwrap_model(model=self.model).config, "patch_size", 2), 1))
return self._prepare_image_crepa_self_flow_batch(batch, state, patch_size=patch_size)
@@ -739,7 +742,7 @@ def controlnet_predict(self, prepared_batch: dict) -> dict:
# Forward pass through the transformer with ControlNet residuals
model_pred = self.model(**hidream_transformer_kwargs)[0]
- return {"model_prediction": model_pred * -1} # the model is trained with inverted velocity :(
+ return {"model_prediction": self.raw_model_prediction_to_model_prediction(model_pred)}
def get_lora_target_layers(self):
manual_targets = self._get_peft_lora_target_modules()
diff --git a/simpletuner/helpers/models/hidream/pipeline.py b/simpletuner/helpers/models/hidream/pipeline.py
index 406e67a7e..0f45de6ca 100644
--- a/simpletuner/helpers/models/hidream/pipeline.py
+++ b/simpletuner/helpers/models/hidream/pipeline.py
@@ -23,7 +23,6 @@
convert_state_dict_to_peft,
convert_unet_state_dict_to_peft,
get_adapter_name,
- get_peft_kwargs,
is_peft_available,
is_peft_version,
is_torch_version,
@@ -46,6 +45,7 @@
)
from simpletuner.helpers.models.hidream.schedule import FlowUniPCMultistepScheduler
+from simpletuner.helpers.training.lora_format import get_peft_kwargs
from simpletuner.helpers.training.lycoris import apply_tlora_inference_mask, clear_tlora_mask
from simpletuner.helpers.utils.offloading import restore_offload_state, unpack_offload_state
diff --git a/simpletuner/helpers/models/ideogram/model.py b/simpletuner/helpers/models/ideogram/model.py
index 9d585a711..0cf2cdfb4 100644
--- a/simpletuner/helpers/models/ideogram/model.py
+++ b/simpletuner/helpers/models/ideogram/model.py
@@ -64,6 +64,9 @@ class Ideogram4(ImageModelFoundation):
def max_swappable_blocks(cls, config=None) -> Optional[int]:
return 33
+ def raw_model_prediction_to_model_prediction(self, raw_prediction: torch.Tensor) -> torch.Tensor:
+ return raw_prediction * -1
+
def setup_model_flavour(self):
super().setup_model_flavour()
flavour = getattr(self.config, "model_flavour", None) or self.DEFAULT_MODEL_FLAVOUR
@@ -662,7 +665,7 @@ def model_predict(self, prepared_batch):
)
packed_prediction = model_output[:, text_tokens:]
model_prediction = self._unpack_latents(packed_prediction, latent_height, latent_width)
- return {"model_prediction": model_prediction * -1}
+ return {"model_prediction": self.raw_model_prediction_to_model_prediction(model_prediction)}
def sample_flow_sigmas(self, batch: dict, state: dict) -> tuple[torch.Tensor, torch.Tensor]:
bsz = batch["latents"].shape[0]
diff --git a/simpletuner/helpers/models/ideogram/pipeline.py b/simpletuner/helpers/models/ideogram/pipeline.py
index de11af4f5..b08bd4d20 100644
--- a/simpletuner/helpers/models/ideogram/pipeline.py
+++ b/simpletuner/helpers/models/ideogram/pipeline.py
@@ -12,7 +12,7 @@
import torch
from diffusers.loaders.lora_base import LoraBaseMixin, _fetch_state_dict
-from diffusers.utils import USE_PEFT_BACKEND, get_adapter_name, get_peft_kwargs, is_peft_version
+from diffusers.utils import USE_PEFT_BACKEND, get_adapter_name, is_peft_version
from huggingface_hub import hf_hub_download
from huggingface_hub.errors import EntryNotFoundError
from PIL import Image
@@ -54,6 +54,7 @@
make_step_intervals,
)
from simpletuner.helpers.models.ideogram.transformer import Ideogram4Config, Ideogram4Transformer
+from simpletuner.helpers.training.lora_format import get_peft_kwargs
def _load_subfolder_state_dict(repo_id: str, subfolder: str, basename: str) -> dict[str, torch.Tensor]:
diff --git a/simpletuner/helpers/models/krea2/lora_pipeline.py b/simpletuner/helpers/models/krea2/lora_pipeline.py
index c1a1528a1..d277ac86f 100644
--- a/simpletuner/helpers/models/krea2/lora_pipeline.py
+++ b/simpletuner/helpers/models/krea2/lora_pipeline.py
@@ -9,7 +9,6 @@
USE_PEFT_BACKEND,
convert_unet_state_dict_to_peft,
get_adapter_name,
- get_peft_kwargs,
is_peft_available,
is_peft_version,
is_torch_version,
@@ -19,6 +18,7 @@
)
from huggingface_hub.utils import validate_hf_hub_args
+from simpletuner.helpers.training.lora_format import get_peft_kwargs
from simpletuner.helpers.utils.offloading import restore_offload_state, unpack_offload_state
_LOW_CPU_MEM_USAGE_DEFAULT_LORA = False
diff --git a/simpletuner/helpers/models/krea2/quantized_loading.py b/simpletuner/helpers/models/krea2/quantized_loading.py
index 1a69c1c3b..8f9106412 100644
--- a/simpletuner/helpers/models/krea2/quantized_loading.py
+++ b/simpletuner/helpers/models/krea2/quantized_loading.py
@@ -128,6 +128,9 @@ def _materialize_krea2_meta_buffers(model: nn.Module) -> None:
def _load_sdnq_training_symbols():
+ from simpletuner.helpers.training.sdnq_compile import configure_sdnq_compile_mode
+
+ configure_sdnq_compile_mode()
try:
from sdnq.dequantizer import SDNQDequantizer
from sdnq.layers import get_sdnq_wrapper_class
diff --git a/simpletuner/helpers/models/longcat_video/pipeline.py b/simpletuner/helpers/models/longcat_video/pipeline.py
index 706541c40..b71266db5 100644
--- a/simpletuner/helpers/models/longcat_video/pipeline.py
+++ b/simpletuner/helpers/models/longcat_video/pipeline.py
@@ -15,7 +15,6 @@
convert_state_dict_to_peft,
convert_unet_state_dict_to_peft,
get_adapter_name,
- get_peft_kwargs,
is_peft_version,
logging,
scale_lora_layers,
diff --git a/simpletuner/helpers/models/ltxvideo2/autoencoder.py b/simpletuner/helpers/models/ltxvideo2/autoencoder.py
index c53141151..262464f65 100644
--- a/simpletuner/helpers/models/ltxvideo2/autoencoder.py
+++ b/simpletuner/helpers/models/ltxvideo2/autoencoder.py
@@ -611,7 +611,7 @@ def forward(
return hidden_states
-# Like LTXVideoUpBlock3d but with no conv_in and the updated LTX2VideoResnetBlock3d
+# Like LTXVideoUpBlock3d but with the updated LTX2VideoResnetBlock3d
class LTX2VideoUpBlock3d(nn.Module):
r"""
Up block used in the LTXVideo model.
@@ -663,10 +663,11 @@ def __init__(
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(in_channels * 4, 0)
self.conv_in = None
- if in_channels != out_channels:
+ upsampler_in_channels = out_channels * upscale_factor
+ if in_channels != upsampler_in_channels:
self.conv_in = LTX2VideoResnetBlock3d(
in_channels=in_channels,
- out_channels=out_channels,
+ out_channels=upsampler_in_channels,
dropout=dropout,
eps=resnet_eps,
non_linearity=resnet_act_fn,
@@ -986,7 +987,7 @@ def __init__(
num_block_out_channels = len(block_out_channels)
self.up_blocks = nn.ModuleList([])
for i in range(num_block_out_channels):
- input_channel = output_channel // upsample_factor[i]
+ input_channel = output_channel
output_channel = block_out_channels[i] // upsample_factor[i]
up_block = LTX2VideoUpBlock3d(
diff --git a/simpletuner/helpers/models/minimaxh3/__init__.py b/simpletuner/helpers/models/minimaxh3/__init__.py
new file mode 100644
index 000000000..a58637272
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/__init__.py
@@ -0,0 +1,21 @@
+from .autoencoder import AutoencoderKLMiniMaxH3
+from .autoencoder_audio import AutoencoderKLMiniMaxH3Audio
+from .modular_blocks_minimax_h3 import MiniMaxH3Blocks, MiniMaxH3Ref2VABlocks
+from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline
+from .pipeline import MiniMaxH3Pipeline
+from .pipeline_ref import MiniMaxH3Ref2VAPipeline
+from .scheduler import MiniMaxH3Scheduler
+from .transformer import MiniMaxH3Transformer3DModel
+
+__all__ = [
+ "AutoencoderKLMiniMaxH3",
+ "AutoencoderKLMiniMaxH3Audio",
+ "MiniMaxH3Blocks",
+ "MiniMaxH3ModularPipeline",
+ "MiniMaxH3Pipeline",
+ "MiniMaxH3Ref2VABlocks",
+ "MiniMaxH3Ref2VAModularPipeline",
+ "MiniMaxH3Ref2VAPipeline",
+ "MiniMaxH3Scheduler",
+ "MiniMaxH3Transformer3DModel",
+]
diff --git a/simpletuner/helpers/models/minimaxh3/activations.py b/simpletuner/helpers/models/minimaxh3/activations.py
new file mode 100644
index 000000000..ea89e9953
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/activations.py
@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MiniMaxH3SwiGLU(nn.Module):
+ def __init__(self, dim_in: int, dim_out: int, bias: bool = True, gate_first: bool = False):
+ super().__init__()
+ self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)
+ self.gate_first = bool(gate_first)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ left, right = self.proj(hidden_states).chunk(2, dim=-1)
+ if self.gate_first:
+ return F.silu(left) * right
+ return left * F.silu(right)
+
+
+class MiniMaxH3FeedForward(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ dim_out: int | None = None,
+ mult: int = 4,
+ inner_dim: int | None = None,
+ dropout: float = 0.0,
+ bias: bool = True,
+ gate_first: bool = False,
+ ):
+ super().__init__()
+ inner_dim = int(dim * mult) if inner_dim is None else inner_dim
+ dim_out = dim if dim_out is None else dim_out
+ self.net = nn.ModuleList(
+ [
+ MiniMaxH3SwiGLU(dim, inner_dim, bias=bias, gate_first=gate_first),
+ nn.Dropout(dropout),
+ nn.Linear(inner_dim, dim_out, bias=bias),
+ ]
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for module in self.net:
+ hidden_states = module(hidden_states)
+ return hidden_states
diff --git a/simpletuner/helpers/models/minimaxh3/autoencoder.py b/simpletuner/helpers/models/minimaxh3/autoencoder.py
new file mode 100644
index 000000000..14f924bfa
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/autoencoder.py
@@ -0,0 +1,1396 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import re
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.autoencoders.vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import logging
+from diffusers.utils.accelerate_utils import apply_forward_hook
+
+from .activations import MiniMaxH3FeedForward
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class MiniMaxH3VideoCausalConv3d(nn.Conv3d):
+ r"""
+ 3D convolution used throughout the MiniMax-H3 video encoder.
+
+ Spatial padding is symmetric and uses `spatial_padding_mode` (`"reflect"` in the released checkpoint); temporal
+ padding is causal, i.e. `kernel_size_t - 1` zero frames are prepended and nothing is appended.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int | tuple[int, int, int],
+ stride: int | tuple[int, int, int] = 1,
+ spatial_padding: int = 0,
+ temporal_padding: int = 0,
+ spatial_padding_mode: str = "reflect",
+ ) -> None:
+ super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=0)
+ self.spatial_padding = spatial_padding
+ self.temporal_padding = temporal_padding
+ self.spatial_padding_mode = spatial_padding_mode
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ if self.spatial_padding > 0:
+ padding = self.spatial_padding
+ hidden_states = F.pad(hidden_states, (padding, padding, padding, padding, 0, 0), mode=self.spatial_padding_mode)
+ if self.temporal_padding > 0:
+ hidden_states = F.pad(hidden_states, (0, 0, 0, 0, self.temporal_padding, 0), mode="constant")
+ return F.conv3d(hidden_states, self.weight, self.bias, stride=self.stride, padding=0, dilation=self.dilation)
+
+
+class MiniMaxH3VideoGroupNorm(nn.GroupNorm):
+ r"""
+ Group normalization applied to each latent frame in isolation (`use_t_isolated_gn` in the original config): the
+ temporal axis is folded into the batch axis so statistics never mix across frames.
+ """
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ batch_size, num_channels, num_frames, height, width = hidden_states.shape
+ hidden_states = hidden_states.permute(0, 2, 1, 3, 4).contiguous()
+ hidden_states = hidden_states.view(batch_size * num_frames, num_channels, 1, height, width)
+ hidden_states = super().forward(hidden_states)
+ hidden_states = hidden_states.view(batch_size, num_frames, num_channels, height, width)
+ return hidden_states.permute(0, 2, 1, 3, 4).contiguous()
+
+
+class MiniMaxH3VideoResnetBlock3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ norm_num_groups: int = 32,
+ norm_eps: float = 1e-6,
+ spatial_padding_mode: str = "reflect",
+ ) -> None:
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ self.norm1 = MiniMaxH3VideoGroupNorm(norm_num_groups, in_channels, eps=norm_eps, affine=True)
+ self.conv1 = MiniMaxH3VideoCausalConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ spatial_padding=1,
+ temporal_padding=2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ self.norm2 = MiniMaxH3VideoGroupNorm(norm_num_groups, out_channels, eps=norm_eps, affine=True)
+ self.conv2 = MiniMaxH3VideoCausalConv3d(
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ spatial_padding=1,
+ temporal_padding=2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ self.conv_shortcut = None
+ if in_channels != out_channels:
+ self.conv_shortcut = MiniMaxH3VideoCausalConv3d(in_channels, out_channels, kernel_size=1)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = F.silu(self.norm1(hidden_states))
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = F.silu(self.norm2(hidden_states))
+ hidden_states = self.conv2(hidden_states)
+ if self.conv_shortcut is not None:
+ residual = self.conv_shortcut(residual)
+ return residual + hidden_states
+
+
+class MiniMaxH3VideoDownsample3d(nn.Module):
+ r"""
+ Strided 3x3x3 downsampling convolution. A spatial stride of 2 is preceded by an asymmetric bottom/right pad of 1
+ (the convolution itself carries no spatial padding), so the output is exactly `ceil(size / 2)`.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ temporal_stride: int = 1,
+ spatial_stride: int = 2,
+ spatial_padding_mode: str = "reflect",
+ ) -> None:
+ super().__init__()
+ self.spatial_stride = spatial_stride
+ self.spatial_padding_mode = spatial_padding_mode
+ self.conv = MiniMaxH3VideoCausalConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=(temporal_stride, spatial_stride, spatial_stride),
+ spatial_padding=0,
+ temporal_padding=2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ if self.spatial_stride == 2:
+ hidden_states = F.pad(hidden_states, (0, 1, 0, 1, 0, 0), mode=self.spatial_padding_mode)
+ return self.conv(hidden_states)
+
+
+class MiniMaxH3VideoDownBlock3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_layers: int,
+ temporal_downsample_factor: int,
+ spatial_downsample_factor: int,
+ norm_num_groups: int = 32,
+ norm_eps: float = 1e-6,
+ spatial_padding_mode: str = "reflect",
+ ) -> None:
+ super().__init__()
+ self.resnets = nn.ModuleList(
+ [
+ MiniMaxH3VideoResnetBlock3d(
+ in_channels=in_channels if i == 0 else out_channels,
+ out_channels=out_channels,
+ norm_num_groups=norm_num_groups,
+ norm_eps=norm_eps,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ for i in range(num_layers)
+ ]
+ )
+ self.downsamplers = None
+ if temporal_downsample_factor * spatial_downsample_factor > 1:
+ self.downsamplers = nn.ModuleList(
+ [
+ MiniMaxH3VideoDownsample3d(
+ out_channels,
+ out_channels,
+ temporal_stride=temporal_downsample_factor,
+ spatial_stride=spatial_downsample_factor,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ ]
+ )
+
+ self.gradient_checkpointing = False
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for resnet in self.resnets:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+ else:
+ hidden_states = resnet(hidden_states)
+ if self.downsamplers is not None:
+ for downsampler in self.downsamplers:
+ hidden_states = downsampler(hidden_states)
+ return hidden_states
+
+
+class MiniMaxH3VideoEncoder3d(nn.Module):
+ r"""
+ Causal 3D CNN encoder. `block_out_channels` gives the channel count of every level; the per-level
+ `spatial_downsample_factors` / `temporal_downsample_factors` multiply out to the total compression ratios.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 48,
+ block_out_channels: tuple[int, ...] = (128, 256, 256, 512, 512, 1024),
+ layers_per_block: int = 2,
+ spatial_downsample_factors: tuple[int, ...] = (2, 2, 2, 2, 1, 1),
+ temporal_downsample_factors: tuple[int, ...] = (1, 2, 2, 1, 1, 1),
+ norm_num_groups: int = 32,
+ norm_eps: float = 1e-6,
+ spatial_padding_mode: str = "reflect",
+ ) -> None:
+ super().__init__()
+
+ self.conv_in = MiniMaxH3VideoCausalConv3d(
+ in_channels,
+ block_out_channels[0],
+ kernel_size=3,
+ spatial_padding=1,
+ temporal_padding=2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ block_in_channels = (block_out_channels[0],) + tuple(block_out_channels[:-1])
+ self.down_blocks = nn.ModuleList(
+ [
+ MiniMaxH3VideoDownBlock3d(
+ in_channels=block_in_channels[i],
+ out_channels=block_out_channels[i],
+ num_layers=layers_per_block,
+ temporal_downsample_factor=temporal_downsample_factors[i],
+ spatial_downsample_factor=spatial_downsample_factors[i],
+ norm_num_groups=norm_num_groups,
+ norm_eps=norm_eps,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ for i in range(len(block_out_channels))
+ ]
+ )
+
+ self.norm_out = MiniMaxH3VideoGroupNorm(norm_num_groups, block_out_channels[-1], eps=norm_eps, affine=True)
+ self.conv_out = MiniMaxH3VideoCausalConv3d(
+ block_out_channels[-1],
+ out_channels,
+ kernel_size=3,
+ spatial_padding=1,
+ temporal_padding=2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.conv_in(hidden_states)
+ for down_block in self.down_blocks:
+ hidden_states = down_block(hidden_states)
+ hidden_states = F.silu(self.norm_out(hidden_states))
+ return self.conv_out(hidden_states)
+
+
+class MiniMaxH3VideoRotaryPosEmbed(nn.Module):
+ r"""
+ 3-axis rotary embedding for the ViT decoder. Coordinates are length-normalized to `[-1, 1)` per axis and scaled by
+ `2 * pi`, and the resulting `(t, h, w)` angles are concatenated and then duplicated, so the first
+ `rope_dim_ratio * attention_head_dim` channels of every head are rotated.
+ """
+
+ def __init__(self, dim: int, theta: float = 100.0, num_axes: int = 3) -> None:
+ super().__init__()
+ if dim % (2 * num_axes) != 0:
+ raise ValueError(f"`dim` {dim} must be divisible by `2 * num_axes` {2 * num_axes}.")
+ inv_freq = 1.0 / theta ** torch.arange(0, 1, 2 * num_axes / dim, dtype=torch.float32)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ angles = 2.0 * math.pi * position_ids[:, :, :, None] * self.inv_freq[None, None, None, :]
+ angles = angles.flatten(2, 3).tile(2).unsqueeze(2)
+ return angles.cos(), angles.sin()
+
+
+class MiniMaxH3VideoAttnProcessor:
+ _attention_backend = None
+ _parallel_config = None
+
+ def __call__(
+ self,
+ attn: "MiniMaxH3VideoAttention",
+ hidden_states: torch.Tensor,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ) -> torch.Tensor:
+ query = attn.to_q(hidden_states).unflatten(2, (attn.heads, -1))
+ key = attn.to_k(hidden_states).unflatten(2, (attn.heads, -1))
+ value = attn.to_v(hidden_states).unflatten(2, (attn.heads, -1))
+
+ # The reference normalizes Q/K in float32 regardless of the compute dtype.
+ query = attn.norm_q(query.float()).to(query.dtype)
+ key = attn.norm_k(key.float()).to(key.dtype)
+
+ if rotary_emb is not None:
+ cos, sin = rotary_emb
+ cos = cos.to(query.dtype)
+ sin = sin.to(query.dtype)
+ rotary_dim = cos.shape[-1]
+ query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
+ key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
+ query_first, query_second = query_rotary.chunk(2, dim=-1)
+ key_first, key_second = key_rotary.chunk(2, dim=-1)
+ query_rotated = torch.cat([-query_second, query_first], dim=-1)
+ key_rotated = torch.cat([-key_second, key_first], dim=-1)
+ query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1)
+ key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1)
+
+ hidden_states = dispatch_attention_fn(
+ query,
+ key,
+ value,
+ attn_mask=None,
+ backend=self._attention_backend,
+ parallel_config=self._parallel_config,
+ )
+ hidden_states = hidden_states.flatten(2, 3)
+ return attn.to_out[0](hidden_states)
+
+
+class MiniMaxH3VideoAttention(nn.Module, AttentionModuleMixin):
+ _default_processor_cls = MiniMaxH3VideoAttnProcessor
+ _available_processors = [MiniMaxH3VideoAttnProcessor]
+
+ def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-5, bias: bool = True) -> None:
+ super().__init__()
+ self.heads = heads
+ self.dim_head = dim_head
+ self.use_bias = bias
+ inner_dim = heads * dim_head
+
+ self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
+ self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
+ self.to_q = nn.Linear(dim, inner_dim, bias=bias)
+ self.to_k = nn.Linear(dim, inner_dim, bias=bias)
+ self.to_v = nn.Linear(dim, inner_dim, bias=bias)
+ self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=bias), nn.Dropout(0.0)])
+
+ self.set_processor(MiniMaxH3VideoAttnProcessor())
+
+ def forward(
+ self, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None
+ ) -> torch.Tensor:
+ return self.processor(self, hidden_states, rotary_emb)
+
+
+class MiniMaxH3VideoTransformerBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ heads: int,
+ dim_head: int,
+ ffn_mult: int = 4,
+ eps: float = 1e-5,
+ bias: bool = True,
+ swiglu_gate_first: bool = True,
+ ) -> None:
+ super().__init__()
+ self.norm1 = nn.RMSNorm(dim, eps=eps, elementwise_affine=True)
+ self.attn = MiniMaxH3VideoAttention(dim=dim, heads=heads, dim_head=dim_head, eps=eps, bias=bias)
+ self.scale1 = nn.Parameter(torch.zeros(dim))
+ self.norm2 = nn.RMSNorm(dim, eps=eps, elementwise_affine=True)
+ self.ff = MiniMaxH3FeedForward(dim, mult=ffn_mult, bias=bias, gate_first=swiglu_gate_first)
+ self.scale2 = nn.Parameter(torch.zeros(dim))
+
+ def forward(
+ self, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None
+ ) -> torch.Tensor:
+ # The reference normalizes in float32 regardless of the compute dtype.
+ norm_hidden_states = self.norm1(hidden_states.float()).to(hidden_states.dtype)
+ scale1 = self.scale1.to(device=hidden_states.device, dtype=hidden_states.dtype)
+ hidden_states = hidden_states + self.attn(norm_hidden_states, rotary_emb) * scale1
+ norm_hidden_states = self.norm2(hidden_states.float()).to(hidden_states.dtype)
+ scale2 = self.scale2.to(device=hidden_states.device, dtype=hidden_states.dtype)
+ hidden_states = hidden_states + self.ff(norm_hidden_states) * scale2
+ return hidden_states
+
+
+class MiniMaxH3VideoViTDecoder3d(nn.Module):
+ r"""
+ Non-causal ViT decoder. Every latent voxel becomes one token; `num_register_tokens` learned register tokens plus a
+ single all-zero token are appended (all at position `0`), attended over with full self-attention, and dropped
+ again before the patch projection expands each token into a `patch_size_t x patch_size x patch_size` pixel block.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 24,
+ out_channels: int = 3,
+ patch_size: int = 16,
+ patch_size_t: int = 4,
+ num_layers: int = 36,
+ num_attention_heads: int = 32,
+ attention_head_dim: int = 64,
+ num_register_tokens: int = 4,
+ ffn_mult: int = 4,
+ swiglu_gate_first: bool = True,
+ rope_theta: float = 100.0,
+ rope_dim_ratio: float = 0.75,
+ norm_eps: float = 1e-5,
+ ) -> None:
+ super().__init__()
+ dim = num_attention_heads * attention_head_dim
+ self.patch_size = patch_size
+ self.patch_size_t = patch_size_t
+ self.out_channels = out_channels
+ self.num_register_tokens = num_register_tokens
+
+ self.rope = MiniMaxH3VideoRotaryPosEmbed(int(attention_head_dim * rope_dim_ratio), theta=rope_theta)
+ self.proj_in = nn.Linear(in_channels, dim)
+ self.register_tokens = nn.Parameter(torch.zeros(1, num_register_tokens, dim))
+ self.transformer_blocks = nn.ModuleList(
+ [
+ MiniMaxH3VideoTransformerBlock(
+ dim=dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ ffn_mult=ffn_mult,
+ eps=norm_eps,
+ swiglu_gate_first=swiglu_gate_first,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+ self.norm_out = nn.LayerNorm(dim, elementwise_affine=True, eps=norm_eps)
+ self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size)
+
+ self.gradient_checkpointing = False
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ batch_size, num_channels, num_frames, height, width = hidden_states.shape
+
+ hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(batch_size, num_frames * height * width, num_channels)
+ hidden_states = self.proj_in(hidden_states)
+ num_patches = hidden_states.shape[1]
+
+ register_tokens = self.register_tokens.to(device=hidden_states.device, dtype=hidden_states.dtype).expand(
+ batch_size, -1, -1
+ )
+ cls_token = torch.zeros_like(hidden_states[:, :1, :])
+ hidden_states = torch.cat([hidden_states, register_tokens, cls_token], dim=1)
+
+ grids = [
+ 2.0 * (torch.arange(0.5, size, dtype=torch.float32, device=hidden_states.device) / size) - 1.0
+ for size in (num_frames, height, width)
+ ]
+ position_ids = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=-1).flatten(0, 2)
+ position_ids = position_ids.unsqueeze(0).expand(batch_size, -1, -1)
+ suffix_ids = position_ids.new_zeros((batch_size, self.num_register_tokens + 1, 3))
+ position_ids = torch.cat([position_ids, suffix_ids], dim=1)
+ rotary_emb = self.rope(position_ids)
+
+ for block in self.transformer_blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(block, hidden_states, rotary_emb)
+ else:
+ hidden_states = block(hidden_states, rotary_emb)
+
+ hidden_states = self.norm_out(hidden_states)
+ hidden_states = self.proj_out(hidden_states)
+ hidden_states = hidden_states[:, :num_patches, :]
+
+ patch_size, patch_size_t = self.patch_size, self.patch_size_t
+ hidden_states = hidden_states.view(
+ batch_size,
+ num_frames,
+ height,
+ width,
+ self.out_channels,
+ patch_size_t,
+ patch_size,
+ patch_size,
+ )
+ hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
+ return hidden_states.reshape(
+ batch_size,
+ self.out_channels,
+ num_frames * patch_size_t,
+ height * patch_size,
+ width * patch_size,
+ )
+
+
+_MINIMAX_H3_VAE_DEFAULT_SPATIAL_DOWNSAMPLE_FACTORS = (2, 2, 2, 2, 1, 1)
+_MINIMAX_H3_VAE_DEFAULT_TEMPORAL_DOWNSAMPLE_FACTORS = (1, 2, 2, 1, 1, 1)
+
+
+def _strip_minimax_h3_vae_checkpoint_prefix(key: str) -> str:
+ for prefix in ("vae.", "first_stage_model.", "model.first_stage_model.", "model.vae."):
+ if key.startswith(prefix):
+ return key[len(prefix) :]
+ return key
+
+
+def _map_minimax_h3_vae_comfy_key_to_diffusers(key: str) -> list[str]:
+ from simpletuner.helpers.models.minimaxh3.transformer import _COMFY_QUANT_METADATA_SUFFIXES
+
+ if key.endswith(_COMFY_QUANT_METADATA_SUFFIXES):
+ return []
+ if key in {"decoder.mask_token", "latents_mean", "latents_std", "pixel_mean", "pixel_std"}:
+ return []
+ if key == "decoder.pos_embed.inv_freq":
+ return ["decoder.rope.inv_freq"]
+
+ block_match = re.match(r"^encoder\.down\.(\d+)\.block\.(\d+)\.(.+)$", key)
+ if block_match is not None:
+ level, block, rest = block_match.groups()
+ rest = rest.replace("nin_shortcut.", "conv_shortcut.", 1)
+ return [f"encoder.down_blocks.{level}.resnets.{block}.{rest}"]
+
+ downsample_match = re.match(r"^encoder\.down\.(\d+)\.downsample\.conv\.(.+)$", key)
+ if downsample_match is not None:
+ level, rest = downsample_match.groups()
+ return [f"encoder.down_blocks.{level}.downsamplers.0.conv.{rest}"]
+
+ key = key.replace("decoder.x_embedder.", "decoder.proj_in.", 1)
+ key = re.sub(r"\.attn\.to_out\.(?!0\.)", ".attn.to_out.0.", key)
+ key = key.replace(".ff.w1.", ".ff.net.0.proj.")
+ key = key.replace(".ff.w2.", ".ff.net.2.")
+ if key.endswith(".attn.to_qkv.weight"):
+ base = key.removesuffix(".attn.to_qkv.weight")
+ return [
+ f"{base}.attn.to_q.weight",
+ f"{base}.attn.to_k.weight",
+ f"{base}.attn.to_v.weight",
+ ]
+ if key.endswith(".attn.to_qkv.bias"):
+ base = key.removesuffix(".attn.to_qkv.bias")
+ return [
+ f"{base}.attn.to_q.bias",
+ f"{base}.attn.to_k.bias",
+ f"{base}.attn.to_v.bias",
+ ]
+ return [key]
+
+
+def _count_minimax_h3_vae_indexed_blocks(keys: set[str], prefix: str) -> int:
+ indices = set()
+ for key in keys:
+ if not key.startswith(prefix):
+ continue
+ index = key[len(prefix) :].split(".", 1)[0]
+ if index.isdigit():
+ indices.add(int(index))
+ return max(indices) + 1 if indices else 0
+
+
+def _get_minimax_h3_vae_checkpoint_tensor(checkpoint, stripped_key: str) -> torch.Tensor:
+ for raw_key in checkpoint.keys():
+ if _strip_minimax_h3_vae_checkpoint_prefix(raw_key) == stripped_key:
+ return checkpoint.get_tensor(raw_key)
+ raise RuntimeError(f"MiniMax-H3 VAE checkpoint is missing required tensor: {stripped_key}")
+
+
+def _infer_minimax_h3_vae_config_from_checkpoint(checkpoint) -> dict[str, Any]:
+ raw_keys = {_strip_minimax_h3_vae_checkpoint_prefix(key) for key in checkpoint.keys()}
+ if "encoder.conv_in.weight" not in raw_keys:
+ raise RuntimeError("MiniMax-H3 VAE single-file checkpoint does not contain recognized encoder keys.")
+
+ conv_in = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "encoder.conv_in.weight")
+ conv_out = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "encoder.conv_out.weight")
+ post_quant = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "post_quant_conv.weight")
+ decoder_proj_in_key = (
+ "decoder.x_embedder.weight" if "decoder.x_embedder.weight" in raw_keys else "decoder.proj_in.weight"
+ )
+ decoder_proj_in = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, decoder_proj_in_key)
+ decoder_proj_out = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "decoder.proj_out.weight")
+
+ block_out_channels = []
+ level_count = _count_minimax_h3_vae_indexed_blocks(raw_keys, "encoder.down.")
+ diffusers_level_count = _count_minimax_h3_vae_indexed_blocks(raw_keys, "encoder.down_blocks.")
+ for index in range(max(level_count, diffusers_level_count)):
+ comfy_key = f"encoder.down.{index}.block.0.conv2.weight"
+ diffusers_key = f"encoder.down_blocks.{index}.resnets.0.conv2.weight"
+ if comfy_key in raw_keys:
+ block_out_channels.append(_get_minimax_h3_vae_checkpoint_tensor(checkpoint, comfy_key).shape[0])
+ elif diffusers_key in raw_keys:
+ block_out_channels.append(_get_minimax_h3_vae_checkpoint_tensor(checkpoint, diffusers_key).shape[0])
+ if not block_out_channels:
+ block_out_channels = [conv_in.shape[0]]
+
+ layers_per_block = _count_minimax_h3_vae_indexed_blocks(raw_keys, "encoder.down.0.block.")
+ if layers_per_block == 0:
+ layers_per_block = _count_minimax_h3_vae_indexed_blocks(raw_keys, "encoder.down_blocks.0.resnets.")
+ layers_per_block = layers_per_block or 2
+
+ num_levels = len(block_out_channels)
+ if num_levels == len(_MINIMAX_H3_VAE_DEFAULT_SPATIAL_DOWNSAMPLE_FACTORS):
+ spatial_downsample_factors = _MINIMAX_H3_VAE_DEFAULT_SPATIAL_DOWNSAMPLE_FACTORS
+ temporal_downsample_factors = _MINIMAX_H3_VAE_DEFAULT_TEMPORAL_DOWNSAMPLE_FACTORS
+ else:
+ spatial_downsample_factors = tuple(
+ (
+ 2
+ if (
+ f"encoder.down.{index}.downsample.conv.weight" in raw_keys
+ or f"encoder.down_blocks.{index}.downsamplers.0.conv.weight" in raw_keys
+ )
+ else 1
+ )
+ for index in range(num_levels)
+ )
+ temporal_downsample_factors = (1,) * num_levels
+
+ decoder_hidden_size = decoder_proj_in.shape[0]
+ if "decoder.transformer_blocks.0.attn.to_qkv.weight" in raw_keys:
+ decoder_qkv = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "decoder.transformer_blocks.0.attn.to_qkv.weight")
+ decoder_hidden_size = decoder_qkv.shape[1]
+ elif "decoder.transformer_blocks.0.attn.to_q.weight" in raw_keys:
+ decoder_hidden_size = _get_minimax_h3_vae_checkpoint_tensor(
+ checkpoint, "decoder.transformer_blocks.0.attn.to_q.weight"
+ ).shape[1]
+
+ if decoder_hidden_size % 64 == 0:
+ decoder_attention_head_dim = 64
+ else:
+ decoder_attention_head_dim = decoder_hidden_size
+ decoder_num_attention_heads = decoder_hidden_size // decoder_attention_head_dim
+
+ has_raw_swiglu = "decoder.transformer_blocks.0.ff.w1.weight" in raw_keys
+ has_diffusers_swiglu = "decoder.transformer_blocks.0.ff.net.0.proj.weight" in raw_keys
+ if has_raw_swiglu and has_diffusers_swiglu:
+ raise RuntimeError("MiniMax-H3 VAE checkpoint mixes raw and Diffusers SwiGLU key layouts.")
+ if has_raw_swiglu:
+ ffn_weight = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "decoder.transformer_blocks.0.ff.w1.weight")
+ else:
+ ffn_weight = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "decoder.transformer_blocks.0.ff.net.0.proj.weight")
+ decoder_ffn_mult = max(ffn_weight.shape[0] // (2 * decoder_hidden_size), 1)
+
+ decoder_num_register_tokens = 4
+ if "decoder.register_tokens" in raw_keys:
+ decoder_num_register_tokens = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, "decoder.register_tokens").shape[1]
+
+ latent_channels = decoder_proj_in.shape[1]
+ norm_num_groups = 32
+ if any(channel % norm_num_groups != 0 for channel in block_out_channels):
+ norm_num_groups = 1
+ inferred_config: dict[str, Any] = {
+ "in_channels": conv_in.shape[1],
+ "out_channels": 3,
+ "latent_channels": latent_channels,
+ "block_out_channels": tuple(block_out_channels),
+ "layers_per_block": layers_per_block,
+ "spatial_downsample_factors": tuple(spatial_downsample_factors),
+ "temporal_downsample_factors": tuple(temporal_downsample_factors),
+ "norm_num_groups": norm_num_groups,
+ "decoder_num_layers": _count_minimax_h3_vae_indexed_blocks(raw_keys, "decoder.transformer_blocks."),
+ "decoder_num_attention_heads": decoder_num_attention_heads,
+ "decoder_attention_head_dim": decoder_attention_head_dim,
+ "decoder_num_register_tokens": decoder_num_register_tokens,
+ "decoder_ffn_mult": decoder_ffn_mult,
+ # The official Diffusers conversion swaps raw `[gate; up]` tensors to `[up; gate]`.
+ "decoder_swiglu_gate_first": has_raw_swiglu,
+ }
+ if tuple(conv_out.shape[:2]) != (2 * latent_channels, block_out_channels[-1]):
+ raise RuntimeError(
+ "MiniMax-H3 VAE checkpoint uses unsupported encoder/decoder latent channel wiring: "
+ f"encoder.conv_out has shape {tuple(conv_out.shape)}, decoder.proj_in has {latent_channels} channels."
+ )
+ if tuple(post_quant.shape[:2]) != (latent_channels, latent_channels):
+ raise RuntimeError(
+ "MiniMax-H3 VAE checkpoint uses unsupported post_quant_conv shape "
+ f"{tuple(post_quant.shape)}; SimpleTuner expects matching embed/z channels."
+ )
+ if inferred_config["decoder_num_layers"] <= 0:
+ raise RuntimeError("MiniMax-H3 VAE checkpoint does not contain decoder transformer blocks.")
+
+ for config_key in ("latents_mean", "latents_std"):
+ if config_key not in raw_keys:
+ continue
+ value = _get_minimax_h3_vae_checkpoint_tensor(checkpoint, config_key).to(torch.float32).flatten()
+ if value.shape[0] != latent_channels:
+ raise RuntimeError(
+ f"MiniMax-H3 VAE checkpoint {config_key} has {value.shape[0]} values, expected {latent_channels}."
+ )
+ inferred_config[config_key] = tuple(float(item) for item in value.tolist())
+ inferred_config.setdefault("latents_mean", (0.0,) * latent_channels)
+ inferred_config.setdefault("latents_std", (1.0,) * latent_channels)
+
+ # Comfy fixes the decoder output channels at RGB. Allow explicit test/config overrides for nonstandard fixtures.
+ patch_volume = math.prod(inferred_config["spatial_downsample_factors"]) ** 2 * math.prod(
+ inferred_config["temporal_downsample_factors"]
+ )
+ if decoder_proj_out.shape[0] % patch_volume == 0:
+ inferred_config["out_channels"] = decoder_proj_out.shape[0] // patch_volume
+ return inferred_config
+
+
+def _set_minimax_h3_vae_module_buffer(root: nn.Module, buffer_name: str, value: torch.Tensor) -> None:
+ module = root
+ parts = buffer_name.split(".")
+ for part in parts[:-1]:
+ module = getattr(module, part)
+ module._buffers[parts[-1]] = value
+
+
+def _normalize_minimax_h3_vae_convrot_scale(scale: torch.Tensor, out_features: int, key: str) -> torch.Tensor:
+ scale = scale.to(torch.float32)
+ if scale.ndim == 0 or scale.numel() == 1:
+ return scale.reshape(1, 1).expand(out_features, 1).contiguous()
+ if tuple(scale.shape) == (out_features,):
+ return scale.reshape(out_features, 1).contiguous()
+ if tuple(scale.shape) == (out_features, 1):
+ return scale.contiguous()
+ raise RuntimeError(
+ f"MiniMax-H3 VAE ConvRot tensor {key}_scale has shape {tuple(scale.shape)}, " f"expected {(out_features, 1)}."
+ )
+
+
+def _split_minimax_h3_vae_comfy_qkv_tensor(
+ tensor: torch.Tensor,
+ mapped_keys: list[str],
+ expected_state_dict: dict[str, torch.Tensor],
+ head_dim: int,
+ key: str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ if len(mapped_keys) != 3:
+ raise RuntimeError(f"MiniMax-H3 VAE fused QKV tensor {key} maps to {len(mapped_keys)} targets, expected 3.")
+ q_shape = expected_state_dict[mapped_keys[0]].shape
+ inner_dim = q_shape[0]
+ if inner_dim % head_dim != 0:
+ raise RuntimeError(
+ f"MiniMax-H3 VAE fused QKV tensor {key} has inner dim {inner_dim}, "
+ f"which is not divisible by decoder_attention_head_dim={head_dim}."
+ )
+ if tensor.shape[0] != inner_dim * 3:
+ raise RuntimeError(
+ f"MiniMax-H3 VAE fused QKV tensor {key} has first dimension {tensor.shape[0]}, " f"expected {inner_dim * 3}."
+ )
+ heads = inner_dim // head_dim
+ qkv = tensor.reshape(heads, 3, head_dim, *tensor.shape[1:])
+ return tuple(part.reshape(inner_dim, *tensor.shape[1:]).contiguous() for part in qkv.unbind(dim=1))
+
+
+class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, AutoencoderMixin):
+ r"""
+ A VAE model with a causal 3D CNN encoder and a non-causal ViT decoder, used in
+ [MiniMax-H3](https://huggingface.co/MiniMaxAI).
+
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+ for all models (such as downloading or saving).
+
+ Latents are normalized with per-channel `latents_mean` / `latents_std` rather than a `scaling_factor`; a pipeline
+ encodes with `(latent - latents_mean) / latents_std` and decodes with `latent * latents_std + latents_mean`.
+
+ The pixel convention is ImageNet-normalized RGB over a `[0, 1]` base range, not the usual `[-1, 1]`: `encode`
+ expects `(pixel - imagenet_mean) / imagenet_std` and `decode` returns values in that same space, so a pipeline has
+ to apply `sample * imagenet_std + imagenet_mean` (mean `(0.485, 0.456, 0.406)`, std `(0.229, 0.224, 0.225)`) and
+ clamp to `[0, 1]` before postprocessing.
+
+ The temporal geometry is fixed by `clip_length` (17 pixel frames per encoder chunk) and `token_drop` (3 trailing
+ latent frames dropped per encode): `17 * n + 5` pixel frames map to `5 * n + 2` latent frames. A single pixel frame
+ is true image mode and maps to one latent frame.
+
+ Unlike most autoencoders in the library, spatial tiling is **on by default**: MiniMax-H3 was released with tiling
+ enabled for both encoding and decoding, and the released frames are the blended-tile ones, so disabling tiling
+ changes the output. Use `enable_tiling` to change the tile geometry, `disable_tiling` to turn it off.
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["MiniMaxH3VideoResnetBlock3d", "MiniMaxH3VideoTransformerBlock"]
+ _repeated_blocks = ["MiniMaxH3VideoTransformerBlock"]
+ _skip_layerwise_casting_patterns = ["norm"]
+ # The released checkpoint is float32 and the verified decode recipe is float16 *autocast over float32 weights*
+ # (see `decode`). A pipeline-level `torch_dtype=torch.bfloat16` must therefore not downcast the weights, so every
+ # top-level module is pinned, mirroring the transformer's mixed-precision contract.
+ _keep_in_fp32_modules = ["encoder", "decoder", "quant_conv", "post_quant_conv"]
+
+ @register_to_config
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ latent_channels: int = 24,
+ block_out_channels: tuple[int, ...] = (128, 256, 256, 512, 512, 1024),
+ layers_per_block: int = 2,
+ spatial_downsample_factors: tuple[int, ...] = (2, 2, 2, 2, 1, 1),
+ temporal_downsample_factors: tuple[int, ...] = (1, 2, 2, 1, 1, 1),
+ norm_num_groups: int = 32,
+ norm_eps: float = 1e-6,
+ spatial_padding_mode: str = "reflect",
+ decoder_num_layers: int = 36,
+ decoder_num_attention_heads: int = 32,
+ decoder_attention_head_dim: int = 64,
+ decoder_num_register_tokens: int = 4,
+ decoder_ffn_mult: int = 4,
+ decoder_swiglu_gate_first: bool = False,
+ decoder_rope_theta: float = 100.0,
+ decoder_rope_dim_ratio: float = 0.75,
+ decoder_norm_eps: float = 1e-5,
+ clip_length: int = 17,
+ token_drop: int = 3,
+ latents_mean: tuple[float, ...] = (0.0,) * 24,
+ latents_std: tuple[float, ...] = (1.0,) * 24,
+ ) -> None:
+ super().__init__()
+
+ self.spatial_compression_ratio = math.prod(spatial_downsample_factors)
+ self.temporal_compression_ratio = math.prod(temporal_downsample_factors)
+
+ self.encoder = MiniMaxH3VideoEncoder3d(
+ in_channels=in_channels,
+ out_channels=2 * latent_channels,
+ block_out_channels=block_out_channels,
+ layers_per_block=layers_per_block,
+ spatial_downsample_factors=spatial_downsample_factors,
+ temporal_downsample_factors=temporal_downsample_factors,
+ norm_num_groups=norm_num_groups,
+ norm_eps=norm_eps,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
+ self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
+ self.decoder = MiniMaxH3VideoViTDecoder3d(
+ in_channels=latent_channels,
+ out_channels=out_channels,
+ patch_size=self.spatial_compression_ratio,
+ patch_size_t=self.temporal_compression_ratio,
+ num_layers=decoder_num_layers,
+ num_attention_heads=decoder_num_attention_heads,
+ attention_head_dim=decoder_attention_head_dim,
+ num_register_tokens=decoder_num_register_tokens,
+ ffn_mult=decoder_ffn_mult,
+ swiglu_gate_first=decoder_swiglu_gate_first,
+ rope_theta=decoder_rope_theta,
+ rope_dim_ratio=decoder_rope_dim_ratio,
+ norm_eps=decoder_norm_eps,
+ )
+
+ # Derived temporal-chunking geometry. `clip_length` pixel frames are encoded at a time; because
+ # `clip_length` is not a multiple of `temporal_compression_ratio`, the decoder has to re-derive the
+ # implicit leading pad (`frame_pre_padding`) and the overlap that `token_drop` leaves behind.
+ self.frame_pre_padding = (-clip_length) % self.temporal_compression_ratio
+ self.tokens_chunk_size = math.ceil(clip_length / self.temporal_compression_ratio)
+ self.token_overlap = (-token_drop) % self.tokens_chunk_size
+ self.frame_overlap = max(self.token_overlap * self.temporal_compression_ratio - self.frame_pre_padding, 0)
+ self.use_temporal_chunking = True
+
+ # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
+ # to perform decoding of a single video latent at a time.
+ self.use_slicing = False
+
+ # When encoding/decoding spatially large videos, the memory requirement is very high. By splitting the frames
+ # into smaller tiles, running the encoder/decoder per tile and blending the overlaps, the memory requirement
+ # can be lowered. MiniMax-H3 ships with tiling enabled.
+ self.use_tiling = True
+
+ # The tile size in pixel space, and the minimum overlap between two neighbouring tiles. The actual overlaps are
+ # widened (in multiples of `spatial_compression_ratio`) so that the tiles cover the frame exactly.
+ self.tile_sample_min_height = 256
+ self.tile_sample_min_width = 256
+ self.tile_sample_min_overlap_height = 64
+ self.tile_sample_min_overlap_width = 64
+
+ @classmethod
+ def from_single_file(
+ cls,
+ pretrained_model_link_or_path: str,
+ *args: Any,
+ filename: str | None = None,
+ subfolder: str | None = None,
+ revision: str | None = None,
+ torch_dtype: torch.dtype | None = None,
+ **kwargs: Any,
+ ) -> "AutoencoderKLMiniMaxH3":
+ del args
+ from simpletuner.helpers.models.minimaxh3.transformer import (
+ _COMFY_QUANT_METADATA_SUFFIXES,
+ _open_minimax_h3_single_file,
+ _resolve_minimax_h3_single_file_path,
+ )
+
+ checkpoint_path = _resolve_minimax_h3_single_file_path(
+ pretrained_model_link_or_path,
+ filename=filename,
+ subfolder=subfolder,
+ revision=revision,
+ )
+ non_quantized_state_dict: dict[str, torch.Tensor] = {}
+ quantized_weights: dict[str, tuple[torch.Tensor, torch.Tensor, int]] = {}
+
+ with _open_minimax_h3_single_file(checkpoint_path) as checkpoint:
+ inferred_config = _infer_minimax_h3_vae_config_from_checkpoint(checkpoint)
+ init_config = {**inferred_config, **kwargs}
+ with torch.device("meta"):
+ model = cls(**init_config)
+ expected_state_dict = model.state_dict()
+ checkpoint_keys = set(checkpoint.keys())
+
+ for raw_key in checkpoint.keys():
+ key = _strip_minimax_h3_vae_checkpoint_prefix(raw_key)
+ if key.endswith(_COMFY_QUANT_METADATA_SUFFIXES):
+ continue
+ mapped_keys = _map_minimax_h3_vae_comfy_key_to_diffusers(key)
+ if not mapped_keys:
+ continue
+
+ tensor = checkpoint.get_tensor(raw_key)
+ if tensor.dtype == torch.int8:
+ from simpletuner.helpers.models.z_image.quantized_loading import _decode_comfy_quant
+
+ scale_key = f"{raw_key}_scale"
+ if scale_key not in checkpoint_keys:
+ raise RuntimeError(f"MiniMax-H3 VAE ConvRot tensor {raw_key} is missing weight_scale")
+ quant_key = f"{raw_key.removesuffix('.weight')}.comfy_quant"
+ if quant_key not in checkpoint_keys:
+ raise RuntimeError(f"MiniMax-H3 VAE ConvRot tensor {raw_key} is missing comfy_quant metadata")
+ quant_metadata = _decode_comfy_quant(checkpoint.get_tensor(quant_key))
+ if not quant_metadata.get("convrot", False):
+ raise RuntimeError(f"MiniMax-H3 VAE INT8 tensor {raw_key} is not marked as ConvRot")
+ hadamard_group_size = int(quant_metadata.get("convrot_groupsize", 0))
+ if hadamard_group_size <= 0:
+ raise RuntimeError(f"MiniMax-H3 VAE ConvRot tensor {raw_key} has invalid convrot_groupsize")
+
+ scale = checkpoint.get_tensor(scale_key)
+ if len(mapped_keys) == 3:
+ scale = _normalize_minimax_h3_vae_convrot_scale(scale, tensor.shape[0], raw_key)
+ qkv_tensors = _split_minimax_h3_vae_comfy_qkv_tensor(
+ tensor,
+ mapped_keys,
+ expected_state_dict,
+ int(model.config.decoder_attention_head_dim),
+ raw_key,
+ )
+ qkv_scales = _split_minimax_h3_vae_comfy_qkv_tensor(
+ scale,
+ mapped_keys,
+ expected_state_dict,
+ int(model.config.decoder_attention_head_dim),
+ f"{raw_key}_scale",
+ )
+ for mapped_key, qkv_tensor, qkv_scale in zip(
+ mapped_keys,
+ qkv_tensors,
+ qkv_scales,
+ ):
+ quantized_weights[mapped_key] = (
+ qkv_tensor,
+ qkv_scale,
+ hadamard_group_size,
+ )
+ elif len(mapped_keys) == 1:
+ scale = _normalize_minimax_h3_vae_convrot_scale(scale, tensor.shape[0], raw_key)
+ quantized_weights[mapped_keys[0]] = (tensor.contiguous(), scale, hadamard_group_size)
+ else:
+ raise RuntimeError(f"MiniMax-H3 VAE ConvRot tensor {raw_key} maps to multiple targets unexpectedly")
+ continue
+
+ if not torch.is_floating_point(tensor):
+ raise RuntimeError(
+ f"MiniMax-H3 VAE tensor {raw_key} has unsupported dtype {tensor.dtype}. "
+ "Only floating-point and INT8 ConvRot single-file VAE tensors are supported."
+ )
+
+ if len(mapped_keys) == 3:
+ qkv_tensors = _split_minimax_h3_vae_comfy_qkv_tensor(
+ tensor,
+ mapped_keys,
+ expected_state_dict,
+ int(model.config.decoder_attention_head_dim),
+ raw_key,
+ )
+ for mapped_key, qkv_tensor in zip(mapped_keys, qkv_tensors):
+ non_quantized_state_dict[mapped_key] = qkv_tensor
+ elif len(mapped_keys) == 1:
+ non_quantized_state_dict[mapped_keys[0]] = tensor
+ else:
+ raise RuntimeError(f"MiniMax-H3 VAE tensor {raw_key} maps to multiple targets unexpectedly")
+
+ decoder_rope_inv_freq = non_quantized_state_dict.pop("decoder.rope.inv_freq", None)
+ if decoder_rope_inv_freq is not None and tuple(decoder_rope_inv_freq.shape) != tuple(
+ model.decoder.rope.inv_freq.shape
+ ):
+ raise RuntimeError(
+ f"MiniMax-H3 VAE tensor decoder.rope.inv_freq has shape {tuple(decoder_rope_inv_freq.shape)}, "
+ f"expected {tuple(model.decoder.rope.inv_freq.shape)}"
+ )
+
+ expected_quantized_keys = set(quantized_weights)
+ keep_fp32_patterns = tuple(getattr(cls, "_keep_in_fp32_modules", ()))
+ for key, tensor in list(non_quantized_state_dict.items()):
+ if key not in expected_state_dict:
+ raise RuntimeError(f"MiniMax-H3 VAE checkpoint has unexpected tensor: {key}")
+ expected_tensor = expected_state_dict[key]
+ if tuple(tensor.shape) != tuple(expected_tensor.shape):
+ raise RuntimeError(
+ f"MiniMax-H3 VAE tensor {key} has shape {tuple(tensor.shape)}, expected {tuple(expected_tensor.shape)}"
+ )
+ if torch_dtype is not None and not any(pattern in key for pattern in keep_fp32_patterns):
+ tensor = tensor.to(torch_dtype)
+ non_quantized_state_dict[key] = tensor
+
+ missing, unexpected = model.load_state_dict(non_quantized_state_dict, strict=False, assign=True)
+ real_missing = [key for key in missing if key not in expected_quantized_keys]
+ if real_missing or unexpected:
+ raise RuntimeError(
+ "MiniMax-H3 VAE checkpoint does not match autoencoder architecture. "
+ f"Missing: {len(real_missing)}, Unexpected: {len(unexpected)}"
+ )
+
+ hadamard_group_sizes: set[int] = set()
+ if quantized_weights:
+ from simpletuner.helpers.models.z_image.quantized_loading import _wrap_convrot_linear
+
+ for weight_key, (weight, scale, hadamard_group_size) in quantized_weights.items():
+ if weight_key not in expected_state_dict:
+ raise RuntimeError(f"MiniMax-H3 VAE ConvRot checkpoint has unexpected tensor: {weight_key}")
+ expected_tensor = expected_state_dict[weight_key]
+ if tuple(weight.shape) != tuple(expected_tensor.shape):
+ raise RuntimeError(
+ f"MiniMax-H3 VAE ConvRot tensor {weight_key} has shape {tuple(weight.shape)}, "
+ f"expected {tuple(expected_tensor.shape)}"
+ )
+ hadamard_group_sizes.add(hadamard_group_size)
+ _wrap_convrot_linear(
+ model,
+ weight_key.removesuffix(".weight"),
+ weight,
+ scale,
+ result_dtype=torch_dtype or torch.bfloat16,
+ hadamard_group_size=hadamard_group_size,
+ )
+ if len(hadamard_group_sizes) != 1:
+ raise RuntimeError(
+ "MiniMax-H3 VAE ConvRot checkpoint uses multiple Hadamard group sizes: "
+ f"{sorted(hadamard_group_sizes)}"
+ )
+ group_size = hadamard_group_sizes.pop()
+ model.quantization_method = "minimax_h3_vae_comfy_convrot_sdnq"
+ model.quantization_config = {
+ "quant_method": "sdnq_training",
+ "weights_dtype": "int8",
+ "quantized_matmul_dtype": "int8",
+ "use_hadamard": True,
+ "hadamard_group_size": group_size,
+ "group_size": -1,
+ "source_format": "comfy_minimax_h3_vae_convrot",
+ }
+
+ if decoder_rope_inv_freq is not None:
+ _set_minimax_h3_vae_module_buffer(model, "decoder.rope.inv_freq", decoder_rope_inv_freq.to(torch.float32))
+ elif model.decoder.rope.inv_freq.is_meta:
+ rope_dim = int(model.config.decoder_attention_head_dim * model.config.decoder_rope_dim_ratio)
+ inv_freq = MiniMaxH3VideoRotaryPosEmbed(
+ rope_dim,
+ theta=getattr(model.config, "decoder_rope_theta", 100.0),
+ ).inv_freq
+ _set_minimax_h3_vae_module_buffer(model, "decoder.rope.inv_freq", inv_freq)
+
+ return model
+
+ def enable_tiling(
+ self,
+ tile_sample_min_height: int | None = None,
+ tile_sample_min_width: int | None = None,
+ tile_sample_min_overlap_height: int | None = None,
+ tile_sample_min_overlap_width: int | None = None,
+ ) -> None:
+ r"""
+ Enable tiled VAE encoding/decoding. When this option is enabled, the VAE splits the frames into tiles, encodes
+ or decodes each tile separately and linearly blends the overlaps back together. This lowers the memory
+ requirement and allows processing larger frames.
+
+ Args:
+ tile_sample_min_height (`int`, *optional*):
+ The tile height in pixel space. Frames taller than this are split along the height dimension.
+ tile_sample_min_width (`int`, *optional*):
+ The tile width in pixel space. Frames wider than this are split along the width dimension.
+ tile_sample_min_overlap_height (`int`, *optional*):
+ The minimum overlap, in pixels, between two consecutive vertical tiles.
+ tile_sample_min_overlap_width (`int`, *optional*):
+ The minimum overlap, in pixels, between two consecutive horizontal tiles.
+ """
+ self.use_tiling = True
+ self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
+ self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+ self.tile_sample_min_overlap_height = tile_sample_min_overlap_height or self.tile_sample_min_overlap_height
+ self.tile_sample_min_overlap_width = tile_sample_min_overlap_width or self.tile_sample_min_overlap_width
+
+ def disable_tiling(self) -> None:
+ self.use_tiling = False
+
+ def enable_slicing(self) -> None:
+ self.use_slicing = True
+
+ def disable_slicing(self) -> None:
+ self.use_slicing = False
+
+ def enable_temporal_chunking(self) -> None:
+ self.use_temporal_chunking = True
+
+ def _split_tiles(self, length: int, tile_size: int, min_overlap: int) -> tuple[list[int], list[int], list[int]]:
+ r"""
+ Lay `tile_size`-wide tiles over `length` pixels. The number of tiles is the smallest one whose union can cover
+ `length` while keeping every overlap at least `min_overlap`; the slack is then distributed round-robin over the
+ overlaps in whole `spatial_compression_ratio` steps so that every tile boundary stays latent-aligned.
+ """
+ if tile_size >= length:
+ return [0], [length], []
+
+ num_tiles = math.ceil(length / tile_size)
+ while tile_size * num_tiles - min_overlap * (num_tiles - 1) - length < 0:
+ num_tiles += 1
+
+ overlaps = [min_overlap] * (num_tiles - 1)
+ remaining = tile_size * num_tiles - sum(overlaps) - length
+ for i in range(remaining // self.spatial_compression_ratio):
+ overlaps[i % (num_tiles - 1)] += self.spatial_compression_ratio
+
+ tile_start_indices = [0]
+ for i in range(num_tiles - 1):
+ tile_start_indices.append(tile_start_indices[-1] + tile_size - overlaps[i])
+ return tile_start_indices, [tile_size] * num_tiles, overlaps
+
+ def _blend(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int) -> torch.Tensor:
+ blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)
+ positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype)
+ shape = [1] * a.ndim
+ shape[dim] = blend_extent
+ weight_a = (1 - positions / blend_extent).view(shape)
+ weight_b = (positions / blend_extent).view(shape)
+
+ slice_a = [slice(None)] * a.ndim
+ slice_a[dim] = slice(-blend_extent, None)
+ slice_b = [slice(None)] * b.ndim
+ slice_b[dim] = slice(0, blend_extent)
+ blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b
+
+ if blend_extent == b.shape[dim]:
+ return blended
+ slice_rest = [slice(None)] * b.ndim
+ slice_rest[dim] = slice(blend_extent, None)
+ return torch.cat([blended, b[tuple(slice_rest)]], dim=dim)
+
+ def _stitch_tiles(
+ self,
+ tiles: list[list[torch.Tensor]],
+ height_overlaps: list[int],
+ width_overlaps: list[int],
+ ) -> torch.Tensor:
+ result_rows = []
+ for i, row in enumerate(tiles):
+ result_row = []
+ for j, tile in enumerate(row):
+ if i > 0:
+ tile = self._blend(tiles[i - 1][j], tile, height_overlaps[i - 1], dim=-2)
+ if j > 0:
+ tile = self._blend(row[j - 1], tile, width_overlaps[j - 1], dim=-1)
+ if i < len(tiles) - 1:
+ tile = tile[..., : -height_overlaps[i], :]
+ if j < len(row) - 1:
+ tile = tile[..., :, : -width_overlaps[j]]
+ result_row.append(tile)
+ result_rows.append(torch.cat(result_row, dim=-1))
+ return torch.cat(result_rows, dim=-2)
+
+ @apply_forward_hook
+ def _encode_clip(self, x: torch.Tensor) -> torch.Tensor:
+ r"""
+ Encode one temporal clip, spatially tiled when tiling is enabled.
+
+ MiniMax-H3 encodes a keyframe or an image reference through this method rather than through [`~encode`],
+ because a single frame must not go through the temporal chunking, so it carries the offload hook too.
+ """
+ if not self.use_tiling:
+ return self.quant_conv(self.encoder(x))
+
+ height, width = x.shape[-2], x.shape[-1]
+ y_indices, y_lengths, y_overlaps = self._split_tiles(
+ height, self.tile_sample_min_height, self.tile_sample_min_overlap_height
+ )
+ x_indices, x_lengths, x_overlaps = self._split_tiles(
+ width, self.tile_sample_min_width, self.tile_sample_min_overlap_width
+ )
+
+ rows = []
+ for i_pos, i_len in zip(y_indices, y_lengths):
+ row = []
+ for j_pos, j_len in zip(x_indices, x_lengths):
+ tile = x[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
+ row.append(self.quant_conv(self.encoder(tile)))
+ rows.append(row)
+
+ latent_y_overlaps = [overlap // self.spatial_compression_ratio for overlap in y_overlaps]
+ latent_x_overlaps = [overlap // self.spatial_compression_ratio for overlap in x_overlaps]
+ return self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps)
+
+ def _decode_clip(self, z: torch.Tensor) -> torch.Tensor:
+ r"""Decode one temporal clip, spatially tiled when tiling is enabled."""
+ if not self.use_tiling:
+ return self.decoder(self.post_quant_conv(z))
+
+ # Tiles are laid out in pixel space and then mapped back onto the latent grid.
+ height = z.shape[-2] * self.spatial_compression_ratio
+ width = z.shape[-1] * self.spatial_compression_ratio
+ y_indices, y_lengths, y_overlaps = self._split_tiles(
+ height, self.tile_sample_min_height, self.tile_sample_min_overlap_height
+ )
+ x_indices, x_lengths, x_overlaps = self._split_tiles(
+ width, self.tile_sample_min_width, self.tile_sample_min_overlap_width
+ )
+
+ ratio = self.spatial_compression_ratio
+ rows = []
+ for i_pos, i_len in zip(y_indices, y_lengths):
+ row = []
+ for j_pos, j_len in zip(x_indices, x_lengths):
+ tile = z[
+ ...,
+ i_pos // ratio : i_pos // ratio + i_len // ratio,
+ j_pos // ratio : j_pos // ratio + j_len // ratio,
+ ]
+ row.append(self.decoder(self.post_quant_conv(tile)))
+ rows.append(row)
+
+ return self._stitch_tiles(rows, y_overlaps, x_overlaps)
+
+ @apply_forward_hook
+ def _encode(self, x: torch.Tensor) -> torch.Tensor:
+ r"""
+ Encode a video in `clip_length`-frame chunks and drop the `token_drop` trailing latent frames.
+
+ MiniMax-H3 encodes a video reference through this method rather than through [`~encode`], because the
+ posterior is sampled under a fixed generator rather than through the distribution object, so it carries the
+ offload hook too.
+ """
+ clip_length = self.config.clip_length
+ num_frames = x.shape[2]
+ if num_frames % clip_length != 0:
+ pad_frames = x[:, :, -1:].repeat(1, 1, (-num_frames) % clip_length, 1, 1)
+ x = torch.cat([x, pad_frames], dim=2)
+
+ moments = torch.cat(
+ [self._encode_clip(x[:, :, i * clip_length : (i + 1) * clip_length]) for i in range(x.shape[2] // clip_length)],
+ dim=2,
+ )
+ if self.config.token_drop > 0:
+ moments = moments[:, :, : -self.config.token_drop]
+ return moments
+
+ def _encode_image_or_video(self, x: torch.Tensor) -> torch.Tensor:
+ if x.shape[2] == 1:
+ return self._encode_clip(x)[:, :, -1:, :, :]
+ return self._encode(x)
+
+ def _decode(self, z: torch.Tensor) -> torch.Tensor:
+ r"""
+ Decode a latent video, mirroring the chunking that `_encode` applied.
+
+ `token_drop` removed the tail of every encoded chunk, so consecutive decoded chunks overlap by
+ `frame_overlap` pixel frames and are linearly cross-faded. Latent frames are repeated at the end when the
+ length is not a whole number of chunks; the extra pixel frames are cut off again at the end.
+ """
+ if z.shape[2] == 1:
+ return self._decode_clip(z)[:, :, -1:, :, :]
+
+ tokens_chunk_size = self.tokens_chunk_size
+ token_drop = self.config.token_drop
+ temporal_ratio = self.temporal_compression_ratio
+ chunk_num_frames = tokens_chunk_size * temporal_ratio
+
+ num_tokens = z.shape[2] + token_drop
+ pad_tokens = (-num_tokens) % tokens_chunk_size
+ num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
+ if pad_tokens > 0:
+ z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
+
+ decoded_chunks = []
+ overlap = None
+ for i in range(num_chunks):
+ start = i * tokens_chunk_size
+ clip = self._decode_clip(z[:, :, start : start + tokens_chunk_size + self.token_overlap])
+ for j in range(int(token_drop > 0) + 1):
+ frame_start = j * chunk_num_frames
+ chunk = clip[:, :, frame_start : frame_start + chunk_num_frames]
+ chunk = chunk[:, :, self.frame_pre_padding :]
+ if j == 0:
+ if overlap is not None:
+ chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
+ decoded_chunks.append(chunk)
+ else:
+ overlap = chunk
+ if overlap is not None:
+ decoded_chunks.append(overlap)
+
+ dec = torch.cat(decoded_chunks, dim=2)
+
+ # `pad_tokens` repeated latent frames produced trailing pixel frames that were never requested. A chunk's
+ # last latent frame only covers `clip_length % temporal_ratio` pixel frames, the others cover `temporal_ratio`.
+ if pad_tokens > 0:
+ intra_tail = self.config.clip_length % temporal_ratio
+ num_tokens_before_pad = z.shape[2] - pad_tokens
+ pad_frames = sum(
+ intra_tail if intra_tail and (num_tokens_before_pad + k) % tokens_chunk_size == 0 else temporal_ratio
+ for k in range(pad_tokens)
+ )
+ dec = dec[:, :, :-pad_frames]
+ return dec
+
+ @apply_forward_hook
+ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput | tuple[torch.Tensor]:
+ r"""
+ Encode a batch of videos into latents.
+
+ Args:
+ x (`torch.Tensor`):
+ Input batch of videos, shape `(batch_size, in_channels, num_frames, height, width)`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.autoencoders.autoencoder_kl.AutoencoderKLOutput`] instead of a plain
+ tuple.
+
+ Returns:
+ The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the encoded latents with
+ `latents_mean` / `latents_std` afterwards.
+ """
+ if self.use_slicing and x.shape[0] > 1:
+ moments = torch.cat([self._encode_image_or_video(x_slice) for x_slice in x.split(1)])
+ else:
+ moments = self._encode_image_or_video(x)
+ posterior = DiagonalGaussianDistribution(moments)
+ if not return_dict:
+ return (posterior,)
+ return AutoencoderKLOutput(latent_dist=posterior)
+
+ @apply_forward_hook
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+ r"""
+ Decode a batch of latent videos.
+
+ Args:
+ z (`torch.Tensor`):
+ Input batch of latent videos, shape `(batch_size, latent_channels, num_latent_frames, height, width)`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
+ The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`.
+ """
+ if self.use_slicing and z.shape[0] > 1:
+ decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)])
+ else:
+ decoded = self._decode(z)
+ if not return_dict:
+ return (decoded,)
+ return DecoderOutput(sample=decoded)
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ sample_posterior: bool = False,
+ generator: torch.Generator | None = None,
+ return_dict: bool = True,
+ ) -> DecoderOutput | tuple[torch.Tensor]:
+ r"""
+ Encode then decode a batch of videos.
+
+ Args:
+ sample (`torch.Tensor`):
+ Input batch of videos, shape `(batch_size, in_channels, num_frames, height, width)`.
+ sample_posterior (`bool`, *optional*, defaults to `False`):
+ Whether to sample the posterior instead of taking its mode.
+ generator (`torch.Generator`, *optional*):
+ Generator used when `sample_posterior=True`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
+ The round-tripped videos, shape `(batch_size, out_channels, num_frames, height, width)`.
+ """
+ posterior = self.encode(sample).latent_dist
+ z = posterior.sample(generator=generator) if sample_posterior else posterior.mode()
+ return self.decode(z, return_dict=return_dict)
diff --git a/simpletuner/helpers/models/minimaxh3/autoencoder_audio.py b/simpletuner/helpers/models/minimaxh3/autoencoder_audio.py
new file mode 100644
index 000000000..8e18727d9
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/autoencoder_audio.py
@@ -0,0 +1,671 @@
+# Copyright 2025 The MiniMax authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MiniMax-H3 audio autoencoder.
+
+Waveform in / waveform out — there is no mel front-end and no separate vocoder:
+
+* the **encoder** is a DAC-lineage strided convolutional stack (Snake activations, weight-normed
+ `Conv1d`) that downsamples by `prod(encoder_rates) = 800`, i.e. 40 latents/s at 32 kHz;
+* a **causal-attention projection** (`pre_block`) rewires the 2048-wide encoder trunk to the
+ 32-channel latent width, followed by the `mean_proj` / `logs_proj` posterior heads;
+* the **decoder** is BigVGAN (anti-aliased SnakeBeta activations, transposed-conv upsamplers, AMP
+ residual blocks) preceded by `dec_in_proj`, upsampling by `prod(decoder_rates) = 800`.
+
+The autoencoder is **mono**. MiniMax-H3 carries stereo as two *batch* items — the pipeline decodes
+`[2, 32, T]` into `[2, 1, samples]` and interleaves at the output boundary — so no stereo handling
+belongs here.
+
+Latents are normalized with per-channel `latents_mean` / `latents_std` (32 floats each) rather than a
+scalar `scaling_factor`; both live in the config and are applied by the pipeline.
+
+Module and parameter names are identical to the original checkpoint, so conversion is a passthrough.
+That includes `torch.nn.utils.weight_norm` (the `weight_g` / `weight_v` spelling, as used by the
+other diffusers audio autoencoders) and the registered Kaiser-window resampling `filter` buffers of
+the anti-aliased activations.
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.autoencoders.vae import DecoderOutput
+from diffusers.models.modeling_utils import ModelMixin, get_parameter_dtype
+from diffusers.utils import BaseOutput
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from diffusers.utils.torch_utils import randn_tensor
+from torch.nn.utils import weight_norm
+
+
+class MiniMaxH3AudioDiagonalGaussianDistribution:
+ r"""Posterior of the MiniMax-H3 audio autoencoder, parameterized as `(mean, log_std)`.
+
+ The checkpoint keeps two separate `Conv1d` heads (`mean_proj`, `logs_proj`) instead of one fused
+ moments projection, and the second head predicts the **log standard deviation**, not the log
+ variance. The two tensors are therefore stored as produced, and `mode()` is bit-for-bit
+ `mean_proj`'s output.
+
+ Args:
+ mean (`torch.Tensor`): Posterior mean, `[batch_size, latent_channels, num_frames]`.
+ logs (`torch.Tensor`): Posterior log standard deviation, same shape as `mean`.
+ """
+
+ def __init__(self, mean: torch.Tensor, logs: torch.Tensor):
+ self.mean = mean
+ self.logs = logs
+ self.std = torch.exp(logs)
+
+ def mode(self) -> torch.Tensor:
+ return self.mean
+
+ def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
+ noise = randn_tensor(self.mean.shape, generator=generator, device=self.mean.device, dtype=self.mean.dtype)
+ return self.mean + self.std * noise
+
+
+@dataclass
+class MiniMaxH3AudioEncoderOutput(BaseOutput):
+ r"""
+ Output of [`AutoencoderKLMiniMaxH3Audio.encode`].
+
+ Args:
+ latent_dist (`MiniMaxH3AudioDiagonalGaussianDistribution`):
+ Posterior over the audio latents. MiniMax-H3 always consumes `latent_dist.mode()`.
+ """
+
+ latent_dist: MiniMaxH3AudioDiagonalGaussianDistribution
+
+
+def _wn_conv1d(*args, **kwargs) -> nn.Module:
+ return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+ r"""Kaiser-windowed sinc low-pass filter of shape `[1, 1, kernel_size]`.
+
+ Kept arithmetically identical to the `alias-free-torch` implementation the checkpoint was trained
+ with, because the resulting tensor is stored as a persistent buffer.
+ """
+ half_size = kernel_size // 2
+
+ attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95
+ if attenuation > 50.0:
+ beta = 0.1102 * (attenuation - 8.7)
+ elif attenuation >= 21.0:
+ beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
+ else:
+ beta = 0.0
+ window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+ if kernel_size % 2 == 0:
+ time = torch.arange(-half_size, half_size) + 0.5
+ else:
+ time = torch.arange(kernel_size) - half_size
+
+ filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
+ # Normalize to sum 1 so a constant input does not leak through the resampler.
+ filter_ /= filter_.sum()
+ return filter_.view(1, 1, kernel_size)
+
+
+class MiniMaxH3AudioSnake1d(nn.Module):
+ r"""`x + (alpha + 1e-9)^-1 * sin(alpha * x)^2` over `[batch_size, channels, length]`, with a
+ per-channel learnable `alpha` of shape `[1, channels, 1]`. Used throughout the DAC encoder."""
+
+ def __init__(self, channels: int):
+ super().__init__()
+ self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ alpha = self.alpha.to(device=hidden_states.device, dtype=hidden_states.dtype)
+ return hidden_states + (alpha + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
+
+
+class MiniMaxH3AudioSnakeBeta(nn.Module):
+ r"""`x + (exp(beta) + 1e-9)^-1 * sin(exp(alpha) * x)^2` over `[batch_size, channels, length]`.
+
+ The BigVGAN decoder's activation: separate frequency (`alpha`) and magnitude (`beta`) parameters,
+ both stored in log space as `[channels]` vectors.
+ """
+
+ def __init__(self, channels: int):
+ super().__init__()
+ self.alpha = nn.Parameter(torch.zeros(channels))
+ self.beta = nn.Parameter(torch.zeros(channels))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ alpha = torch.exp(self.alpha.to(device=hidden_states.device, dtype=hidden_states.dtype).unsqueeze(0).unsqueeze(-1))
+ beta = torch.exp(self.beta.to(device=hidden_states.device, dtype=hidden_states.dtype).unsqueeze(0).unsqueeze(-1))
+ return hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
+
+
+class MiniMaxH3AudioLowPassFilter1d(nn.Module):
+ r"""Depthwise Kaiser-sinc low-pass filter with a stride, i.e. the anti-aliased downsampler."""
+
+ def __init__(self, cutoff: float, half_width: float, stride: int, kernel_size: int):
+ super().__init__()
+ even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(even)
+ self.pad_right = kernel_size // 2
+ self.stride = stride
+ self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ num_channels = hidden_states.shape[1]
+ hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate")
+ return F.conv1d(hidden_states, self.filter.expand(num_channels, -1, -1), stride=self.stride, groups=num_channels)
+
+
+class MiniMaxH3AudioUpSample1d(nn.Module):
+ r"""Anti-aliased `ratio`x upsampler (transposed depthwise Kaiser-sinc convolution)."""
+
+ def __init__(self, ratio: int, kernel_size: int):
+ super().__init__()
+ self.ratio = ratio
+ self.stride = ratio
+ self.pad = kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (kernel_size - self.stride) // 2
+ self.pad_right = self.pad * self.stride + (kernel_size - self.stride + 1) // 2
+ self.register_buffer(
+ "filter",
+ kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size),
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ num_channels = hidden_states.shape[1]
+ hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate")
+ hidden_states = self.ratio * F.conv_transpose1d(
+ hidden_states, self.filter.expand(num_channels, -1, -1), stride=self.stride, groups=num_channels
+ )
+ return hidden_states[..., self.pad_left : -self.pad_right]
+
+
+class MiniMaxH3AudioDownSample1d(nn.Module):
+ r"""Anti-aliased `ratio`x downsampler."""
+
+ def __init__(self, ratio: int, kernel_size: int):
+ super().__init__()
+ self.lowpass = MiniMaxH3AudioLowPassFilter1d(
+ cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return self.lowpass(hidden_states)
+
+
+class MiniMaxH3AudioActivation1d(nn.Module):
+ r"""Upsample -> activation -> downsample: the alias-free activation wrapper used by BigVGAN."""
+
+ def __init__(self, activation: nn.Module, ratio: int = 2, kernel_size: int = 12):
+ super().__init__()
+ self.act = activation
+ self.upsample = MiniMaxH3AudioUpSample1d(ratio, kernel_size)
+ self.downsample = MiniMaxH3AudioDownSample1d(ratio, kernel_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.upsample(hidden_states)
+ hidden_states = self.act(hidden_states)
+ return self.downsample(hidden_states)
+
+
+class MiniMaxH3AudioResidualUnit(nn.Module):
+ r"""DAC residual unit: `Snake -> dilated Conv1d(k=7) -> Snake -> Conv1d(k=1)`, plus a shortcut
+ that is center-cropped when the dilated convolution shrinks the time axis."""
+
+ def __init__(self, dim: int, dilation: int):
+ super().__init__()
+ self.block = nn.Sequential(
+ MiniMaxH3AudioSnake1d(dim),
+ _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=((7 - 1) * dilation) // 2),
+ MiniMaxH3AudioSnake1d(dim),
+ _wn_conv1d(dim, dim, kernel_size=1),
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ residual = self.block(hidden_states)
+ pad = (hidden_states.shape[-1] - residual.shape[-1]) // 2
+ if pad > 0:
+ hidden_states = hidden_states[..., pad:-pad]
+ return hidden_states + residual
+
+
+class MiniMaxH3AudioEncoderBlock(nn.Module):
+ r"""Three residual units at dilations 1/3/9, then a strided channel-doubling convolution."""
+
+ def __init__(self, dim: int, stride: int):
+ super().__init__()
+ self.block = nn.Sequential(
+ MiniMaxH3AudioResidualUnit(dim // 2, dilation=1),
+ MiniMaxH3AudioResidualUnit(dim // 2, dilation=3),
+ MiniMaxH3AudioResidualUnit(dim // 2, dilation=9),
+ MiniMaxH3AudioSnake1d(dim // 2),
+ _wn_conv1d(
+ dim // 2,
+ dim,
+ kernel_size=2 * stride,
+ stride=stride,
+ padding=math.ceil(stride / 2),
+ ),
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return self.block(hidden_states)
+
+
+class MiniMaxH3AudioEncoder(nn.Module):
+ r"""DAC waveform encoder: `[batch_size, 1, samples] -> [batch_size, latent_dim, samples / 800]`."""
+
+ def __init__(self, d_model: int, strides: tuple[int, ...], d_latent: int):
+ super().__init__()
+ block: list[nn.Module] = [_wn_conv1d(1, d_model, kernel_size=7, padding=3)]
+ for stride in strides:
+ d_model *= 2
+ block.append(MiniMaxH3AudioEncoderBlock(d_model, stride=stride))
+ block += [
+ MiniMaxH3AudioSnake1d(d_model),
+ _wn_conv1d(d_model, d_latent, kernel_size=3, padding=1),
+ ]
+ self.block = nn.Sequential(*block)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return self.block(hidden_states)
+
+
+class MiniMaxH3AudioGeGluMlp(nn.Module):
+ r"""Pre-norm GeGLU MLP used inside the attention projection block."""
+
+ def __init__(self, in_features: int, hidden_features: int):
+ super().__init__()
+ self.norm = nn.LayerNorm(in_features)
+ self.act = nn.GELU(approximate="tanh")
+ self.w0 = nn.Linear(in_features, hidden_features)
+ self.w1 = nn.Linear(in_features, hidden_features)
+ self.w2 = nn.Linear(hidden_features, in_features)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.act(self.w0(hidden_states)) * self.w1(hidden_states)
+ return self.w2(hidden_states)
+
+
+class MiniMaxH3AudioAttnProcessor:
+ r"""Processor of [`MiniMaxH3AudioCausalAttention`].
+
+ The causal mask is expressed as `is_causal=True` rather than as a materialized mask. Every
+ attention backend honours that flag, with two exceptions: `_native_npu`, whose kernel takes no
+ causal argument and would compute *bidirectional* attention, and context parallelism, which
+ raises for causal attention.
+ """
+
+ _attention_backend = None
+ _parallel_config = None
+
+ def __call__(self, attn: "MiniMaxH3AudioCausalAttention", hidden_states: torch.Tensor) -> torch.Tensor:
+ batch_size, seq_len, _ = hidden_states.shape
+ qkv = F.linear(
+ input=hidden_states,
+ weight=attn.qkv.weight,
+ bias=torch.cat((attn.q_bias, attn.zero_k_bias, attn.v_bias)).to(
+ device=hidden_states.device, dtype=hidden_states.dtype
+ ),
+ )
+ query, key, value = (
+ qkv.reshape(batch_size, seq_len, 3, attn.num_heads, attn.head_dim).permute(2, 0, 1, 3, 4).unbind(0)
+ )
+ hidden_states = dispatch_attention_fn(
+ query,
+ key,
+ value,
+ attn_mask=None,
+ is_causal=True,
+ backend=self._attention_backend,
+ parallel_config=self._parallel_config,
+ )
+ # The heads are mean-pooled away instead of being concatenated, and the head dimension that
+ # remains is adaptively average-pooled down to `out_dim`.
+ hidden_states = torch.mean(hidden_states, dim=2)
+ hidden_states = F.adaptive_avg_pool1d(hidden_states, attn.out_dim)
+ return attn.proj(hidden_states)
+
+
+class MiniMaxH3AudioCausalAttention(nn.Module, AttentionModuleMixin):
+ r"""Causal self-attention that narrows the feature width from `in_dim` to `out_dim`.
+
+ QKV is a single bias-less `nn.Linear`; query and value biases are separate parameters and the key
+ bias is a frozen zero buffer (`zero_k_bias`), exactly as stored in the checkpoint. Heads are
+ `in_dim // num_heads` wide; instead of being concatenated they are **mean-pooled away**, and the
+ remaining head dimension is adaptively average-pooled down to `out_dim`.
+ """
+
+ _default_processor_cls = MiniMaxH3AudioAttnProcessor
+ _available_processors = [MiniMaxH3AudioAttnProcessor]
+ # The checkpoint stores one fused `qkv` projection, so there is nothing to fuse.
+ _supports_qkv_fusion = False
+
+ def __init__(self, in_dim: int, out_dim: int, num_heads: int):
+ super().__init__()
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.head_dim = in_dim // num_heads
+ self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False)
+ self.q_bias = nn.Parameter(torch.zeros(in_dim))
+ self.v_bias = nn.Parameter(torch.zeros(in_dim))
+ self.register_buffer("zero_k_bias", torch.zeros(in_dim))
+ self.proj = nn.Linear(out_dim, out_dim)
+
+ self.set_processor(MiniMaxH3AudioAttnProcessor())
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return self.processor(self, hidden_states)
+
+
+class MiniMaxH3AudioAttnProjection(nn.Module):
+ r"""`pre_block`: residual causal-attention + GeGLU block that rewires `latent_dim` -> `latent_channels`."""
+
+ def __init__(self, in_dim: int, out_dim: int, num_heads: int, mlp_ratio: int = 2):
+ super().__init__()
+ self.norm1 = nn.LayerNorm(in_dim)
+ self.attn = MiniMaxH3AudioCausalAttention(in_dim, out_dim, num_heads)
+ self.proj = nn.Linear(in_dim, out_dim)
+ self.norm3 = nn.LayerNorm(in_dim)
+ self.norm2 = nn.LayerNorm(out_dim)
+ self.mlp = MiniMaxH3AudioGeGluMlp(in_features=out_dim, hidden_features=out_dim * mlp_ratio)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.proj(self.norm3(hidden_states)) + self.attn(self.norm1(hidden_states))
+ return hidden_states + self.mlp(self.norm2(hidden_states))
+
+
+class MiniMaxH3AudioAMPBlock(nn.Module):
+ r"""BigVGAN anti-aliased multi-periodicity block (`AMPBlock1`).
+
+ Each dilation contributes a `(dilated conv, dilation-1 conv)` pair, and every convolution is
+ preceded by its own alias-free SnakeBeta activation.
+ """
+
+ def __init__(self, channels: int, kernel_size: int, dilation: tuple[int, ...]):
+ super().__init__()
+ self.convs1 = nn.ModuleList(
+ [_wn_conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2) for d in dilation]
+ )
+ self.convs2 = nn.ModuleList(
+ [_wn_conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2) for _ in dilation]
+ )
+ self.activations = nn.ModuleList(
+ [MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels)) for _ in range(2 * len(dilation))]
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ acts1, acts2 = self.activations[::2], self.activations[1::2]
+ for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
+ residual = conv1(act1(hidden_states))
+ residual = conv2(act2(residual))
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+class MiniMaxH3AudioBigVGANDecoder(nn.Module):
+ r"""BigVGAN decoder: `[batch_size, latent_dim, num_frames] -> [batch_size, 1, num_frames * 800]`."""
+
+ def __init__(
+ self,
+ in_channels: int,
+ upsample_initial_channel: int,
+ upsample_rates: tuple[int, ...],
+ upsample_kernel_sizes: tuple[int, ...],
+ resblock_kernel_sizes: tuple[int, ...],
+ resblock_dilation_sizes: tuple[tuple[int, ...], ...],
+ ):
+ super().__init__()
+ self.num_kernels = len(resblock_kernel_sizes)
+ self.num_upsamples = len(upsample_rates)
+
+ self.conv_pre = _wn_conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
+
+ # Each upsampler is wrapped in a one-element `ModuleList` in the original checkpoint
+ # (`ups..0`); the extra nesting is kept so the state dict stays a passthrough.
+ self.ups = nn.ModuleList()
+ for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+ self.ups.append(
+ nn.ModuleList(
+ [
+ weight_norm(
+ nn.ConvTranspose1d(
+ upsample_initial_channel // (2**i),
+ upsample_initial_channel // (2 ** (i + 1)),
+ kernel,
+ rate,
+ padding=(kernel - rate) // 2,
+ )
+ )
+ ]
+ )
+ )
+
+ self.resblocks = nn.ModuleList()
+ for i in range(self.num_upsamples):
+ channels = upsample_initial_channel // (2 ** (i + 1))
+ for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+ self.resblocks.append(MiniMaxH3AudioAMPBlock(channels, kernel, tuple(dilation)))
+
+ self.activation_post = MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels))
+ self.conv_post = _wn_conv1d(channels, 1, 7, 1, padding=3, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.conv_pre(hidden_states)
+
+ for i in range(self.num_upsamples):
+ hidden_states = self.ups[i][0](hidden_states)
+ residual = None
+ for j in range(self.num_kernels):
+ block = self.resblocks[i * self.num_kernels + j](hidden_states)
+ residual = block if residual is None else residual + block
+ hidden_states = residual / self.num_kernels
+
+ hidden_states = self.activation_post(hidden_states)
+ hidden_states = self.conv_post(hidden_states)
+ return torch.clamp(hidden_states, min=-1.0, max=1.0)
+
+
+class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin):
+ r"""
+ The audio autoencoder used by [MiniMax-H3](https://huggingface.co/MiniMaxAI): a DAC-lineage
+ convolutional encoder and a BigVGAN decoder, operating directly on mono 32 kHz waveforms.
+
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library
+ implements for all models (such as downloading or saving).
+
+ Args:
+ encoder_dim (`int`, defaults to `64`):
+ Channel width of the encoder's first convolution; doubles at every downsampling stage.
+ encoder_rates (`tuple[int]`, defaults to `(2, 4, 4, 5, 5)`):
+ Encoder strides. Their product (`800`) is the hop length, i.e. 40 latents/s at 32 kHz.
+ latent_dim (`int`, defaults to `2048`):
+ Width of the encoder trunk and of the decoder input, before/after the latent projections.
+ latent_channels (`int`, defaults to `32`):
+ Width of the diffusion latent, i.e. the `mean_proj` / `logs_proj` output channels.
+ num_attention_heads (`int`, defaults to `8`):
+ Number of heads in the causal-attention projection `pre_block`.
+ decoder_dim (`int`, defaults to `1024`):
+ BigVGAN initial channel count; halved at every upsampling stage.
+ decoder_rates (`tuple[int]`, defaults to `(5, 5, 2, 2, 2, 2, 2)`):
+ BigVGAN upsampling rates. Their product must equal `prod(encoder_rates)`.
+ decoder_kernel_sizes (`tuple[int]`, defaults to `(9, 9, 4, 4, 4, 4, 4)`):
+ Transposed-convolution kernel size per upsampling stage.
+ resblock_kernel_sizes (`tuple[int]`, defaults to `(3, 7, 11)`):
+ Kernel sizes of the parallel AMP residual blocks at each upsampling stage.
+ resblock_dilation_sizes (`tuple[tuple[int]]`, defaults to `((1, 3, 5), (1, 3, 5), (1, 3, 5))`):
+ Per-AMP-block dilations.
+ sampling_rate (`int`, defaults to `32000`):
+ Waveform sampling rate.
+ latents_mean (`list[float]`, *optional*):
+ Per-channel latent mean the pipeline uses to normalize / denormalize latents.
+ latents_std (`list[float]`, *optional*):
+ Per-channel latent standard deviation the pipeline uses to normalize / denormalize latents.
+ """
+
+ _supports_gradient_checkpointing = False
+ # The released checkpoint is float32 and the DAC/BigVGAN stack (weight-normalized convolutions, Snake
+ # activations) degrades audibly under bfloat16 (roughly 20 dB quieter decodes), so a pipeline-level
+ # `torch_dtype=torch.bfloat16` must not downcast the weights.
+ _keep_in_fp32_modules = ["encoder", "decoder", "pre_block", "dec_in_proj", "mean_proj", "logs_proj"]
+
+ @register_to_config
+ def __init__(
+ self,
+ encoder_dim: int = 64,
+ encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5),
+ latent_dim: int = 2048,
+ latent_channels: int = 32,
+ num_attention_heads: int = 8,
+ decoder_dim: int = 1024,
+ decoder_rates: tuple[int, ...] = (5, 5, 2, 2, 2, 2, 2),
+ decoder_kernel_sizes: tuple[int, ...] = (9, 9, 4, 4, 4, 4, 4),
+ resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11),
+ resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
+ sampling_rate: int = 32000,
+ latents_mean: list[float] | None = None,
+ latents_std: list[float] | None = None,
+ ):
+ super().__init__()
+
+ encoder_rates = tuple(int(rate) for rate in encoder_rates)
+ decoder_rates = tuple(int(rate) for rate in decoder_rates)
+ self.hop_length = math.prod(encoder_rates)
+ if math.prod(decoder_rates) != self.hop_length:
+ raise ValueError(
+ f"`decoder_rates` must upsample by the encoder hop length {self.hop_length}, got "
+ f"{math.prod(decoder_rates)}."
+ )
+ if latent_dim % latent_channels != 0:
+ raise ValueError(f"`latent_dim` ({latent_dim}) must be a multiple of `latent_channels` ({latent_channels}).")
+
+ self.encoder = MiniMaxH3AudioEncoder(d_model=encoder_dim, strides=encoder_rates, d_latent=latent_dim)
+ self.pre_block = MiniMaxH3AudioAttnProjection(latent_dim, latent_channels, num_heads=num_attention_heads)
+ self.mean_proj = nn.Conv1d(latent_channels, latent_channels, 1)
+ self.logs_proj = nn.Conv1d(latent_channels, latent_channels, 1)
+
+ self.dec_in_proj = nn.Conv1d(latent_channels, latent_dim, 1)
+ self.decoder = MiniMaxH3AudioBigVGANDecoder(
+ in_channels=latent_dim,
+ upsample_initial_channel=decoder_dim,
+ upsample_rates=decoder_rates,
+ upsample_kernel_sizes=tuple(int(kernel) for kernel in decoder_kernel_sizes),
+ resblock_kernel_sizes=tuple(int(kernel) for kernel in resblock_kernel_sizes),
+ resblock_dilation_sizes=tuple(tuple(int(d) for d in dilation) for dilation in resblock_dilation_sizes),
+ )
+
+ @apply_forward_hook
+ def encode(
+ self, sample: torch.Tensor, return_dict: bool = True
+ ) -> MiniMaxH3AudioEncoderOutput | tuple[MiniMaxH3AudioDiagonalGaussianDistribution]:
+ r"""
+ Encode a waveform into the audio latent posterior.
+
+ The waveform is right-padded to a multiple of `hop_length` (800 samples) first. MiniMax-H3
+ always consumes the posterior **mean** (`latent_dist.mode()`) — the `logs_proj` head is never
+ evaluated by the reference pipeline.
+
+ Args:
+ sample (`torch.Tensor`):
+ Mono waveform of shape `[batch_size, 1, samples]`. MiniMax-H3 passes the two stereo
+ channels of a reference clip as `batch_size = 2`.
+ return_dict (`bool`, defaults to `True`):
+ Whether to return a [`MiniMaxH3AudioEncoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`MiniMaxH3AudioEncoderOutput`] or `tuple`:
+ The latent posterior over `[batch_size, latent_channels, samples / 800]`.
+ """
+ if sample.ndim != 3 or sample.shape[1] != 1:
+ raise ValueError(f"`sample` must have shape [batch_size, 1, samples], got {tuple(sample.shape)}.")
+
+ right_pad = math.ceil(sample.shape[-1] / self.hop_length) * self.hop_length - sample.shape[-1]
+ if right_pad > 0:
+ sample = F.pad(sample, (0, right_pad))
+
+ encoder_dtype = get_parameter_dtype(self.encoder)
+ hidden_states = self.encoder(sample.to(encoder_dtype))
+ hidden_states = self.pre_block(hidden_states.transpose(1, 2)).transpose(1, 2)
+ mean, logs = self.mean_proj(hidden_states), self.logs_proj(hidden_states)
+ if encoder_dtype != torch.float32:
+ mean, logs = mean.float(), logs.float()
+
+ posterior = MiniMaxH3AudioDiagonalGaussianDistribution(mean, logs)
+ if not return_dict:
+ return (posterior,)
+ return MiniMaxH3AudioEncoderOutput(latent_dist=posterior)
+
+ @apply_forward_hook
+ def decode(self, latents: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+ r"""
+ Decode audio latents into a waveform.
+
+ Args:
+ latents (`torch.Tensor`):
+ Denormalized latents of shape `[batch_size, latent_channels, num_frames]`. MiniMax-H3
+ passes the two stereo channels as `batch_size = 2`.
+ return_dict (`bool`, defaults to `True`):
+ Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
+ Waveform of shape `[batch_size, 1, num_frames * 800]`, clamped to `[-1, 1]`.
+ """
+ if latents.ndim != 3:
+ raise ValueError(
+ f"`latents` must have shape [batch_size, latent_channels, num_frames], got {tuple(latents.shape)}."
+ )
+
+ decoder_dtype = get_parameter_dtype(self.decoder)
+ decoded = self.decoder(self.dec_in_proj(latents.to(decoder_dtype)))
+ if decoder_dtype != torch.float32:
+ decoded = decoded.float()
+
+ if not return_dict:
+ return (decoded,)
+ return DecoderOutput(sample=decoded)
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ sample_posterior: bool = False,
+ return_dict: bool = True,
+ generator: torch.Generator | None = None,
+ ) -> DecoderOutput | tuple[torch.Tensor]:
+ r"""
+ Encode then decode a waveform.
+
+ Args:
+ sample (`torch.Tensor`):
+ Mono waveform of shape `[batch_size, 1, samples]`.
+ sample_posterior (`bool`, defaults to `False`):
+ Whether to sample the posterior instead of taking its mode. MiniMax-H3 uses the mode.
+ return_dict (`bool`, defaults to `True`):
+ Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple.
+ generator (`torch.Generator`, *optional*):
+ Generator used when `sample_posterior=True`.
+
+ Returns:
+ [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
+ The round-tripped waveform of shape `[batch_size, 1, num_frames * 800]`, clamped to `[-1, 1]`.
+ """
+ posterior = self.encode(sample).latent_dist
+ latents = posterior.sample(generator=generator) if sample_posterior else posterior.mode()
+ return self.decode(latents, return_dict=return_dict)
diff --git a/simpletuner/helpers/models/minimaxh3/before_denoise.py b/simpletuner/helpers/models/minimaxh3/before_denoise.py
new file mode 100644
index 000000000..bc5cb31e1
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/before_denoise.py
@@ -0,0 +1,495 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+
+from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline
+from .packing import (
+ MINIMAX_H3_AUDIO_CHANNELS,
+ MINIMAX_H3_KEYFRAME_NOISE_AUG,
+ MiniMaxH3PackedSequence,
+ build_packed_sequence,
+ build_row_timesteps,
+ patchify_video_latents,
+)
+from .packing_ref2va import MiniMaxH3PreparedReference, build_ref2va_packed_sequence
+from .scheduler import MiniMaxH3Scheduler
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _layout_inputs() -> list[InputParam]:
+ r"""What both packed layouts are built from, beyond the conditioning of the task itself."""
+ return [
+ InputParam(
+ name="text_token_tags",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The per-row modality tag of every row of `prompt_embeds`.",
+ ),
+ InputParam(name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames."),
+ InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."),
+ InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."),
+ InputParam(
+ name="num_audio_latents",
+ type_hint=int,
+ required=True,
+ description="Number of audio latents per channel.",
+ ),
+ ]
+
+
+def _layout_outputs() -> list[OutputParam]:
+ r"""The row layout of the packed sequence, shared by the two tasks."""
+ return [
+ OutputParam(
+ "layout",
+ type_hint=MiniMaxH3PackedSequence,
+ description="The structural description of the packed sequence.",
+ ),
+ OutputParam(
+ "position_ids",
+ type_hint=torch.Tensor,
+ description="The `(t, h, w)` rotary coordinate of every row, in float64.",
+ ),
+ OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."),
+ OutputParam(
+ "video_indices",
+ type_hint=torch.Tensor,
+ description="Sequence positions of the video rows, conditioning rows first.",
+ ),
+ OutputParam(
+ "audio_indices",
+ type_hint=torch.Tensor,
+ description="Sequence positions of the audio rows, reference rows first.",
+ ),
+ OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."),
+ OutputParam(
+ "num_condition_video_rows",
+ type_hint=int,
+ description="How many leading video rows are conditioning rows rather than generated rows.",
+ ),
+ OutputParam(
+ "num_condition_audio_rows",
+ type_hint=int,
+ description="How many leading audio rows are reference rows rather than generated rows.",
+ ),
+ ]
+
+
+def _single_text_token_tags(text_token_tags: torch.Tensor, name: str) -> torch.Tensor:
+ if text_token_tags.ndim == 2 and text_token_tags.shape[0] == 1:
+ text_token_tags = text_token_tags[0]
+ if text_token_tags.ndim != 1:
+ raise ValueError(f"MiniMax-H3 `{name}` must be 1-D, or batched with one row, got {tuple(text_token_tags.shape)}.")
+ return text_token_tags
+
+
+def _set_layout_state(block_state, layout: MiniMaxH3PackedSequence, device: torch.device, prefix: str = "") -> None:
+ setattr(block_state, f"{prefix}layout", layout)
+ setattr(block_state, f"{prefix}position_ids", layout.position_ids.to(device))
+ setattr(block_state, f"{prefix}token_tags", layout.token_tags.to(device))
+ setattr(block_state, f"{prefix}video_indices", layout.video_indices.to(device))
+ setattr(block_state, f"{prefix}audio_indices", layout.audio_indices.to(device))
+ setattr(block_state, f"{prefix}text_indices", layout.text_indices.to(device))
+ if not prefix:
+ block_state.num_condition_video_rows = layout.num_condition_video_rows
+ block_state.num_condition_audio_rows = layout.num_condition_audio_rows
+
+
+class MiniMaxH3PrepareLayoutStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Builds the packed layout of a `t2va` / `fl2va` request — `[text | keyframe conditions | target audio | "
+ "target video]` — and its fp64 rotary grid. MiniMax-H3 runs full self-attention over this one sequence, "
+ "so the layout is what every later block addresses rows through."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ *_layout_inputs(),
+ InputParam(
+ name="negative_text_token_tags",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch text tags for real CFG.",
+ ),
+ InputParam(
+ name="keyframe_anchors",
+ type_hint=tuple,
+ default=(),
+ description="Which end of the video every keyframe is anchored to, in packed order.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return _layout_outputs()
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ text_token_tags = _single_text_token_tags(block_state.text_token_tags, "text_token_tags")
+
+ layout = build_packed_sequence(
+ text_token_tags,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components.patch_size,
+ block_state.keyframe_anchors,
+ )
+ _set_layout_state(block_state, layout, components._execution_device)
+ negative_text_token_tags = getattr(block_state, "negative_text_token_tags", None)
+ if negative_text_token_tags is not None:
+ negative_text_token_tags = _single_text_token_tags(negative_text_token_tags, "negative_text_token_tags")
+ negative_layout = build_packed_sequence(
+ negative_text_token_tags,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components.patch_size,
+ block_state.keyframe_anchors,
+ )
+ _set_layout_state(block_state, negative_layout, components._execution_device, prefix="negative_")
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VAPrepareLayoutStep(ModularPipelineBlocks):
+ model_name = "minimax-h3-ref2va"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Builds the packed layout of a `ref2va` request — `[text | reference blocks | target audio | target "
+ "video]` — and its fp64 rotary grid. The reference order advances the shared audio/video rotary clock, so "
+ "it is part of the layout rather than a detail of the presentation."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ *_layout_inputs(),
+ InputParam(
+ name="negative_text_token_tags",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch text tags for real CFG.",
+ ),
+ InputParam(
+ name="prepared_references",
+ type_hint=list[MiniMaxH3PreparedReference],
+ required=True,
+ description="The prepared references, in packed order, with their latent geometry filled in.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return _layout_outputs()
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ text_token_tags = _single_text_token_tags(block_state.text_token_tags, "text_token_tags")
+
+ layout = build_ref2va_packed_sequence(
+ text_token_tags,
+ block_state.prepared_references,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components.patch_size,
+ )
+ _set_layout_state(block_state, layout, components._execution_device)
+ negative_text_token_tags = getattr(block_state, "negative_text_token_tags", None)
+ if negative_text_token_tags is not None:
+ negative_text_token_tags = _single_text_token_tags(negative_text_token_tags, "negative_text_token_tags")
+ negative_layout = build_ref2va_packed_sequence(
+ negative_text_token_tags,
+ block_state.prepared_references,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components.patch_size,
+ )
+ _set_layout_state(block_state, negative_layout, components._execution_device, prefix="negative_")
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3PrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Draws the initial noise of the generated rows and prepends the conditioning rows. MiniMax-H3 draws the "
+ "video noise as a latent tensor and patchifies it afterwards, then the audio noise directly in row "
+ "layout — both off the request's generator, after the conditioning noise of the encoder step."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames."),
+ InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."),
+ InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."),
+ InputParam(
+ name="num_audio_latents",
+ type_hint=int,
+ required=True,
+ description="Number of audio latents per channel.",
+ ),
+ InputParam.template(
+ "generator",
+ description=("The generator of the request. The video noise is drawn from it first, then the audio noise."),
+ ),
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description=(
+ "Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, "
+ "used instead of the draw."
+ ),
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ description="Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.",
+ ),
+ InputParam(
+ name="condition_latents",
+ type_hint=torch.Tensor,
+ description="The video conditioning rows to prepend, or None for a request that has none.",
+ ),
+ InputParam(
+ name="audio_condition_latents",
+ type_hint=torch.Tensor,
+ description="The audio conditioning rows to prepend, or None for a request that has none.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The video rows of the packed sequence, conditioning rows first.",
+ ),
+ OutputParam(
+ "audio_latents",
+ type_hint=torch.Tensor,
+ description="The channel-major audio rows of the packed sequence, reference rows first.",
+ ),
+ ]
+
+ @staticmethod
+ def prepare_latents(
+ components,
+ num_latent_frames: int,
+ latent_height: int,
+ latent_width: int,
+ num_audio_latents: int,
+ device: torch.device,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ audio_latents: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ r"""
+ Draw the initial noise of both modalities and pack it into transformer rows.
+
+ A request draws every stream from the one generator it is given, and the order is part of what that generator
+ reproduces: the conditioning noise of the keyframes or references first (one draw per condition, in
+ [`~modular_pipelines.minimax_h3.packing.keyframe_condition_noise`]), then the video noise here, as a latent tensor
+ that is patchified afterwards, then the audio noise, directly in row layout. Passing `latents` or
+ `audio_latents` skips its draw and shifts the ones after it.
+
+ Args:
+ num_latent_frames (`int`): Number of video latent frames.
+ latent_height (`int`): Latent height.
+ latent_width (`int`): Latent width.
+ num_audio_latents (`int`): Number of audio latents per channel.
+ device (`torch.device`): The device the rows are drawn on.
+ generator (`torch.Generator`, *optional*): The generator of the request.
+ latents (`torch.Tensor`, *optional*):
+ Pre-generated video noise of shape `(1, latent_channels, num_latent_frames, latent_height,
+ latent_width)`, used instead of the draw.
+ audio_latents (`torch.Tensor`, *optional*):
+ Pre-generated audio noise of shape `(2, audio_latent_channels, num_audio_latents)`.
+
+ Returns:
+ `tuple[torch.Tensor, torch.Tensor]`: the video rows and the channel-major audio rows.
+ """
+ if latents is None:
+ latents = randn_tensor(
+ (1, components.vae_latent_channels, num_latent_frames, latent_height, latent_width),
+ generator=generator,
+ device=device,
+ dtype=torch.float32,
+ )
+ video_rows = patchify_video_latents(latents.to(torch.float32), components.patch_size)
+
+ if audio_latents is None:
+ audio_rows = randn_tensor(
+ (num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, components.audio_latent_channels),
+ generator=generator,
+ device=device,
+ dtype=torch.float32,
+ )
+ else:
+ audio_rows = audio_latents.to(torch.float32).permute(0, 2, 1).reshape(-1, components.audio_latent_channels)
+ return video_rows.to(device), audio_rows.to(device)
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latents, audio_latents = self.prepare_latents(
+ components,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components._execution_device,
+ block_state.generator,
+ block_state.latents,
+ block_state.audio_latents,
+ )
+ if block_state.condition_latents is not None:
+ latents = torch.cat([block_state.condition_latents, latents])
+ if block_state.audio_condition_latents is not None:
+ audio_latents = torch.cat([block_state.audio_condition_latents, audio_latents])
+ block_state.latents, block_state.audio_latents = latents, audio_latents
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3SetTimestepsStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Initializes the two schedules — `shift = 12.0` for video, `shift = 3.0` for audio — and stages the "
+ "row-to-timestep plan of every step. One forward serves every modality and every noise level at once: "
+ "the generated rows step down their own schedule while the conditioning rows stay pinned at their "
+ "noise-augmentation level, and that assignment is static per step."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", required=True),
+ InputParam(
+ name="layout",
+ type_hint=MiniMaxH3PackedSequence,
+ required=True,
+ description="The structural description of the packed sequence.",
+ ),
+ InputParam(
+ name="negative_layout",
+ type_hint=MiniMaxH3PackedSequence,
+ description="Optional negative-branch structural description for real CFG.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps of the video schedule."),
+ OutputParam("audio_timesteps", type_hint=torch.Tensor, description="Timesteps of the audio schedule."),
+ OutputParam(
+ "row_timestep_plan",
+ type_hint=list,
+ description=(
+ "One `(timestep, timestep_indices)` pair per step: the distinct timesteps of the sequence and the "
+ "index of every row into them."
+ ),
+ ),
+ OutputParam(
+ "negative_row_timestep_plan",
+ type_hint=list,
+ description="Optional row-to-timestep plan for the negative branch.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+ block_state.timesteps = components.scheduler.timesteps
+ if components.audio_scheduler is not None:
+ components.audio_scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+ block_state.audio_timesteps = components.audio_scheduler.timesteps
+ else:
+ block_state.audio_timesteps = block_state.timesteps
+
+ block_state.row_timestep_plan = [
+ tuple(
+ tensor.to(device)
+ for tensor in build_row_timesteps(
+ block_state.layout,
+ float(timestep),
+ float(audio_timestep),
+ max(float(timestep), MINIMAX_H3_KEYFRAME_NOISE_AUG),
+ 1.0,
+ )
+ )
+ for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
+ ]
+ negative_layout = getattr(block_state, "negative_layout", None)
+ if negative_layout is not None:
+ block_state.negative_row_timestep_plan = [
+ tuple(
+ tensor.to(device)
+ for tensor in build_row_timesteps(
+ negative_layout,
+ float(timestep),
+ float(audio_timestep),
+ max(float(timestep), MINIMAX_H3_KEYFRAME_NOISE_AUG),
+ 1.0,
+ )
+ )
+ for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
+ ]
+ else:
+ block_state.negative_row_timestep_plan = None
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/simpletuner/helpers/models/minimaxh3/before_encoder.py b/simpletuner/helpers/models/minimaxh3/before_encoder.py
new file mode 100644
index 000000000..4831d892f
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/before_encoder.py
@@ -0,0 +1,454 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import PIL
+import torch
+from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+from diffusers.utils import logging
+from PIL import Image, ImageOps
+
+from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline
+from .packing import (
+ MINIMAX_H3_CANVAS_MULTIPLE,
+ MINIMAX_H3_FPS,
+ MINIMAX_H3_MAX_DURATION,
+ MINIMAX_H3_MIN_DURATION,
+ align_num_frames,
+ audio_latent_num_frames,
+ prepare_keyframe_image,
+ resolve_canvas_size,
+ video_latent_num_frames,
+)
+from .packing_ref2va import (
+ MINIMAX_H3_MAX_REFERENCE_AUDIOS,
+ MINIMAX_H3_MAX_REFERENCE_IMAGES,
+ MINIMAX_H3_MAX_REFERENCE_VIDEOS,
+ MINIMAX_H3_MAX_REFERENCES,
+ MiniMaxH3PreparedReference,
+ MiniMaxH3Reference,
+ prepare_reference_frames,
+ prepare_reference_image,
+ prepare_reference_waveform,
+ reference_kind,
+ reference_media_to_uint8,
+ resample_reference_frames,
+ resolve_reference_image_size,
+)
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+MINIMAX_H3_TARGET_MODES = ("auto", "video", "av")
+
+
+def _latent_geometry(components, height: int, width: int, num_frames: int) -> tuple[int, int, int, int]:
+ r"""The latent geometry the packed layout, the noise draws and the decoders all key off."""
+ ratio = components.vae_spatial_compression_ratio
+ return video_latent_num_frames(num_frames), height // ratio, width // ratio, audio_latent_num_frames(num_frames)
+
+
+def _latent_geometry_outputs() -> list[OutputParam]:
+ r"""The declaration of what [`_latent_geometry`] resolves, shared by the two setup blocks."""
+ return [
+ OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."),
+ OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."),
+ OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."),
+ OutputParam("num_audio_latents", type_hint=int, description="Number of generated audio latents per channel."),
+ ]
+
+
+def _target_mode_input() -> InputParam:
+ return InputParam(
+ name="minimax_h3_target_mode",
+ type_hint=str,
+ default="auto",
+ description=(
+ "MiniMax-H3 target modality mode. `auto` and `video` omit target audio rows; `av` keeps joint "
+ "audio-video rows."
+ ),
+ )
+
+
+def _resolve_h3_target_mode(value: str | None) -> str:
+ if value is None or value == "":
+ value = "auto"
+ mode = str(value).strip().lower()
+ if mode not in MINIMAX_H3_TARGET_MODES:
+ raise ValueError(f"`minimax_h3_target_mode` must be one of {', '.join(MINIMAX_H3_TARGET_MODES)}, got {value!r}.")
+ if mode == "auto":
+ return "video"
+ return mode
+
+
+def _apply_h3_target_mode(block_state) -> None:
+ block_state.minimax_h3_target_mode = _resolve_h3_target_mode(getattr(block_state, "minimax_h3_target_mode", "auto"))
+ if block_state.minimax_h3_target_mode == "video":
+ block_state.num_audio_latents = 0
+
+
+class MiniMaxH3SetupStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resolves the plan shared by the `t2va` and `fl2va` tasks: the canvas (MiniMax-H3's own 768-short-edge "
+ "geometry for the aspect ratio of the first keyframe, or 16:9 without keyframes), the image or video "
+ "frame count the VAE can decode, the latent geometry every later block keys off, and the keyframes put "
+ "onto that canvas."
+ )
+
+ @staticmethod
+ def _check_inputs(block_state) -> None:
+ if (block_state.height is None) != (block_state.width is None):
+ raise ValueError("`height` and `width` have to be passed together, or neither of them.")
+ if block_state.height is not None and (
+ block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE
+ ):
+ raise ValueError(
+ f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got "
+ f"{block_state.height}x{block_state.width}."
+ )
+ # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has
+ # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds.
+ aligned_num_frames = align_num_frames(block_state.num_frames)
+ duration = aligned_num_frames / MINIMAX_H3_FPS
+ if aligned_num_frames != 1 and not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION:
+ raise ValueError(
+ f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at "
+ f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can "
+ f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and "
+ f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames} (rounded up to "
+ f"{aligned_num_frames})."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="image",
+ type_hint=PIL.Image.Image,
+ description=(
+ "Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is "
+ "derived from its own aspect ratio."
+ ),
+ ),
+ InputParam(
+ name="last_image",
+ type_hint=PIL.Image.Image,
+ description=(
+ "Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with "
+ "`image` it is the follower of the two and is cover-cropped onto the canvas."
+ ),
+ ),
+ InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."),
+ InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."),
+ InputParam(
+ name="num_frames",
+ type_hint=int,
+ default=124,
+ description=(
+ "Number of frames to generate. `1` renders an image; otherwise the fixed 24 fps video frame count "
+ "is snapped up to the next `17 * n + 5` the video VAE can decode, and the resulting duration must "
+ "stay between 5 and 15 seconds."
+ ),
+ ),
+ _target_mode_input(),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ OutputParam(
+ "num_frames",
+ type_hint=int,
+ description="Resolved number of frames, either 1 or of the form 17 * n + 5.",
+ ),
+ *_latent_geometry_outputs(),
+ OutputParam(
+ "keyframes",
+ type_hint=list,
+ description="The keyframes put onto the target canvas, in packed order (empty for `t2va`).",
+ ),
+ OutputParam(
+ "keyframe_anchors",
+ type_hint=tuple,
+ description="Which end of the video every keyframe is anchored to, in packed order.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self._check_inputs(block_state)
+
+ keyframes = [
+ ImageOps.exif_transpose(keyframe).convert("RGB")
+ for keyframe in (block_state.image, block_state.last_image)
+ if keyframe is not None
+ ]
+ block_state.keyframe_anchors = tuple(
+ anchor
+ for anchor, keyframe in (("first", block_state.image), ("last", block_state.last_image))
+ if keyframe is not None
+ )
+ if block_state.height is None:
+ block_state.height, block_state.width = resolve_canvas_size(*(keyframes[0].size if keyframes else (16, 9)))
+
+ aligned_num_frames = align_num_frames(block_state.num_frames)
+ if aligned_num_frames != block_state.num_frames:
+ logger.warning(
+ f"`num_frames` has to be 1 or of the form 17 * n + 5 for the video VAE; rounding "
+ f"{block_state.num_frames} up to {aligned_num_frames}."
+ )
+ block_state.num_frames = aligned_num_frames
+
+ (
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ ) = _latent_geometry(components, block_state.height, block_state.width, block_state.num_frames)
+ _apply_h3_target_mode(block_state)
+
+ block_state.keyframes = [
+ prepare_keyframe_image(keyframe, block_state.height, block_state.width, stretch=index == 0)
+ for index, keyframe in enumerate(keyframes)
+ ]
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VASetupStep(ModularPipelineBlocks):
+ model_name = "minimax-h3-ref2va"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resolves the `ref2va` plan: the canvas (MiniMax-H3's own 16:9 unless asked otherwise — references never "
+ "bind the generated geometry), the references prepared at their own resolutions, the frame count they "
+ "imply when it was left open, and the latent geometry every later block keys off."
+ )
+
+ @staticmethod
+ def _check_inputs(components, block_state) -> None:
+ if (block_state.height is None) != (block_state.width is None):
+ raise ValueError("`height` and `width` have to be passed together, or neither of them.")
+ if block_state.height is not None and (
+ block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE
+ ):
+ raise ValueError(
+ f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got "
+ f"{block_state.height}x{block_state.width}."
+ )
+ # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has
+ # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds.
+ aligned_num_frames = None if block_state.num_frames is None else align_num_frames(block_state.num_frames)
+ duration = None if aligned_num_frames is None else aligned_num_frames / MINIMAX_H3_FPS
+ if (
+ duration is not None
+ and aligned_num_frames != 1
+ and not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION
+ ):
+ raise ValueError(
+ f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at "
+ f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can "
+ f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and "
+ f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames} (rounded up to "
+ f"{aligned_num_frames})."
+ )
+
+ if not block_state.references:
+ raise ValueError("`ref2va` needs at least one reference; use `MiniMaxH3ModularPipeline` for text-only requests.")
+ kinds = [reference_kind(index, entry) for index, entry in enumerate(block_state.references)]
+ for kind, limit in (
+ ("image", MINIMAX_H3_MAX_REFERENCE_IMAGES),
+ ("video", MINIMAX_H3_MAX_REFERENCE_VIDEOS),
+ ("audio", MINIMAX_H3_MAX_REFERENCE_AUDIOS),
+ ):
+ if kinds.count(kind) > limit:
+ raise ValueError(f"MiniMax-H3 accepts at most {limit} {kind} references, got {kinds.count(kind)}.")
+ if len(kinds) > MINIMAX_H3_MAX_REFERENCES:
+ raise ValueError(
+ f"MiniMax-H3 accepts at most {MINIMAX_H3_MAX_REFERENCES} references in total, got {len(kinds)}."
+ )
+ if set(kinds) == {"audio"}:
+ raise ValueError(
+ "An audio reference has to be paired with at least one image or video reference and cannot be used "
+ "on its own."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="references",
+ type_hint=list[MiniMaxH3Reference],
+ required=True,
+ description=(
+ "The references to condition on, **in the order the model should read them**: the order labels "
+ "them in the prompt presentation and lays them out on the shared rotary clock, so a different "
+ "order is a different request. Every [`MiniMaxH3Reference`] carries exactly one medium, a path or "
+ "in-memory media — `image` (at most 9), `video` at its own `fps` (at most 3, whose `audio` "
+ "soundtrack is conditioned on as well), or `audio` at its own `sample_rate` (at most 3) — for at "
+ "most 12 references in total, and audio references cannot be the only ones. A path is decoded "
+ "when the reference is built, so these blocks only ever see pixels and samples."
+ ),
+ ),
+ InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."),
+ InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."),
+ InputParam(
+ name="num_frames",
+ type_hint=int,
+ description=(
+ "Number of frames to generate. `1` renders an image; otherwise the fixed 24 fps video frame count "
+ "is snapped up to the next `17 * n + 5` the video VAE can decode. May be left out, but only when "
+ "exactly one reference carries audio, in which case the duration is that soundtrack's."
+ ),
+ ),
+ _target_mode_input(),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ OutputParam(
+ "num_frames",
+ type_hint=int,
+ description="Resolved number of frames, either 1 or of the form 17 * n + 5.",
+ ),
+ *_latent_geometry_outputs(),
+ OutputParam(
+ "prepared_references",
+ type_hint=list[MiniMaxH3PreparedReference],
+ description="The references prepared at their own resolutions, in packed order.",
+ ),
+ ]
+
+ @staticmethod
+ def prepare_references(
+ components, references: list[MiniMaxH3Reference], num_frames: int | None
+ ) -> tuple[list[MiniMaxH3PreparedReference], int]:
+ r"""
+ Resolve the references and, if it was left open, the duration they imply.
+
+ Every reference is prepared at its own resolution: an image is resized to a 2048 pixel short edge, a video is
+ resampled onto MiniMax-H3's own 24 fps, rescaled onto the 768 pixel canvas of *its own* aspect ratio and
+ truncated to the generated frame count, and a soundtrack is put on the audio VAE's sample rate and truncated to
+ the generated duration. None of this touches the target canvas.
+
+ A reference that left its `fps` or its `sample_rate` out is taken to already be at MiniMax-H3's own rate, and
+ its frames or its samples then flow through untouched.
+
+ A video reference goes through the two passes the reference implementation's `ffmpeg` decode applied, in the
+ same order: the constant frame rate resample of `resample_reference_frames` and the LANCZOS rescale of
+ `prepare_reference_frames`. Frames handed over at 24 fps and already at the canvas their own aspect ratio
+ resolves to therefore reach the VAE untouched, which is the parity-exact route.
+
+ Args:
+ references (`list[MiniMaxH3Reference]`):
+ The `references` input of a [`MiniMaxH3Ref2VABlocks`] request.
+ num_frames (`int`, *optional*):
+ The requested frame count, or `None` to derive it from the single audio-bearing reference.
+
+ Returns:
+ `tuple[list[MiniMaxH3PreparedReference], int]`: the prepared references, in packed order, and the frame
+ count.
+ """
+ resolved = [
+ MiniMaxH3PreparedReference(kind=reference_kind(index, entry), has_audio=entry.has_audio)
+ for index, entry in enumerate(references)
+ ]
+
+ # The duration may be left open, but then exactly one reference may carry audio, or the request is ambiguous.
+ if num_frames is None:
+ audio_bearing = [index for index, reference in enumerate(resolved) if reference.has_audio]
+ if len(audio_bearing) != 1:
+ raise ValueError(
+ "`num_frames` may only be left to the references when exactly one of them carries audio, got "
+ f"{len(audio_bearing)}."
+ )
+ index = audio_bearing[0]
+ sample_rate = references[index].sample_rate or components.audio_sampling_rate
+ duration = references[index].audio.shape[-1] / sample_rate
+ if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION:
+ raise ValueError(
+ f"`references[{index}]` is {duration:g} seconds long, outside the "
+ f"{MINIMAX_H3_MIN_DURATION} to {MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates."
+ )
+ num_frames = align_num_frames(round(duration * MINIMAX_H3_FPS))
+ # The duration the request generates is the one of the *aligned* frame count, so that is what the
+ # ceiling has to hold for: a 14.99 second soundtrack rounds up to 362 frames, i.e. 15.083 seconds.
+ if num_frames / MINIMAX_H3_FPS > MINIMAX_H3_MAX_DURATION:
+ raise ValueError(
+ f"`references[{index}]` is {duration:g} seconds long, which rounds up to {num_frames} frames "
+ f"(`17 * n + 5`), i.e. {num_frames / MINIMAX_H3_FPS:g} seconds — past the "
+ f"{MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates. Pass `num_frames` to generate a "
+ "shorter video from this soundtrack."
+ )
+ num_frames = align_num_frames(num_frames)
+
+ for reference, entry in zip(resolved, references):
+ if reference.kind == "image":
+ image = entry.image
+ if not isinstance(image, Image.Image):
+ image = Image.fromarray(reference_media_to_uint8(image))
+ image = ImageOps.exif_transpose(image).convert("RGB")
+ height, width = resolve_reference_image_size(*image.size)
+ reference.image = prepare_reference_image(image, height, width)
+ elif reference.kind == "video":
+ frames = resample_reference_frames(reference_media_to_uint8(entry.video), float(entry.fps))
+ reference.frames = prepare_reference_frames(frames, num_frames)
+ if reference.has_audio:
+ reference.waveform = prepare_reference_waveform(
+ entry.audio,
+ entry.sample_rate or components.audio_sampling_rate,
+ components.audio_sampling_rate,
+ max_duration=num_frames / MINIMAX_H3_FPS,
+ )
+ return resolved, num_frames
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self._check_inputs(components, block_state)
+
+ if block_state.height is None:
+ block_state.height, block_state.width = resolve_canvas_size(16, 9)
+
+ requested_num_frames = block_state.num_frames
+ block_state.prepared_references, block_state.num_frames = self.prepare_references(
+ components, block_state.references, block_state.num_frames
+ )
+ if requested_num_frames is not None and requested_num_frames != block_state.num_frames:
+ logger.warning(
+ f"`num_frames` has to be 1 or of the form 17 * n + 5 for the video VAE; rounding "
+ f"{requested_num_frames} up to {block_state.num_frames}."
+ )
+
+ (
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ ) = _latent_geometry(components, block_state.height, block_state.width, block_state.num_frames)
+ _apply_h3_target_mode(block_state)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/simpletuner/helpers/models/minimaxh3/decoders.py b/simpletuner/helpers/models/minimaxh3/decoders.py
new file mode 100644
index 000000000..bd49c0359
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/decoders.py
@@ -0,0 +1,194 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from diffusers.configuration_utils import FrozenDict
+from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.utils import logging
+from diffusers.video_processor import VideoProcessor
+
+from .autoencoder import AutoencoderKLMiniMaxH3
+from .autoencoder_audio import AutoencoderKLMiniMaxH3Audio
+from .modular_pipeline import MiniMaxH3ModularPipeline
+from .packing import MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD, unpack_audio_tokens, unpatchify_video_tokens
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class MiniMaxH3VideoDecodeStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Unpacks the generated video rows back into latents, denormalizes them and decodes them into video. The "
+ "spatial tiling of the video VAE covers the canvas exactly, so the decoded frames need no crop back, but "
+ "the decode itself runs under float16 autocast even though the VAE weights are float32, and the VAE "
+ "produces ImageNet-normalized RGB that is reverted here."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLMiniMaxH3),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ # The video VAE decodes into ImageNet-normalized RGB over a [0, 1] base range, which this block
+ # reverts itself, so the processor must not denormalize a second time.
+ config=FrozenDict({"vae_scale_factor": 16, "do_normalize": False}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The denoised video rows of the packed sequence, conditioning rows first.",
+ ),
+ InputParam(
+ name="num_condition_video_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading video rows are conditioning rows and are dropped here.",
+ ),
+ InputParam(name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames."),
+ InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."),
+ InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."),
+ InputParam.template(
+ "output_type", description="Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("videos", description="The generated video.")]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ latents = unpatchify_video_tokens(
+ block_state.latents[block_state.num_condition_video_rows :],
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ components.vae_latent_channels,
+ components.patch_size,
+ )
+ latents_mean = torch.tensor(components.vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1)
+ latents_std = torch.tensor(components.vae.config.latents_std, device=device).view(1, -1, 1, 1, 1)
+ latents = latents * latents_std + latents_mean
+
+ if block_state.output_type == "latent":
+ block_state.videos = latents
+ else:
+ with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
+ video = components.vae.decode(latents, return_dict=False)[0]
+ pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1)
+ pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1)
+ video = (video.float() * pixel_std + pixel_mean).clamp(0, 1)
+ block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3AudioDecodeStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Unpacks the generated audio rows back into latents, denormalizes them and decodes them into a stereo "
+ "waveform. The audio VAE is mono and takes the two stereo channels as two batch items."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The denoised audio rows of the packed sequence, reference rows first.",
+ ),
+ InputParam(
+ name="num_condition_audio_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading audio rows are reference rows and are dropped here.",
+ ),
+ InputParam(
+ name="num_audio_latents",
+ type_hint=int,
+ required=True,
+ description="Number of audio latents per channel.",
+ ),
+ InputParam.template(
+ "output_type", description="Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "audio",
+ type_hint=torch.Tensor,
+ description="The generated soundtrack, of shape `(1, 2, num_samples)`.",
+ ),
+ OutputParam(
+ "sampling_rate",
+ type_hint=int,
+ description="Sample rate of the generated soundtrack in Hz.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ if int(block_state.num_audio_latents or 0) == 0:
+ block_state.audio = None
+ block_state.sampling_rate = components.audio_sampling_rate
+ self.set_block_state(state, block_state)
+ return components, state
+
+ audio_latents = unpack_audio_tokens(
+ block_state.audio_latents[block_state.num_condition_audio_rows :], block_state.num_audio_latents
+ )
+ audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean, device=device).view(1, -1, 1)
+ audio_latents_std = torch.tensor(components.audio_vae.config.latents_std, device=device).view(1, -1, 1)
+ audio_latents = audio_latents * audio_latents_std + audio_latents_mean
+
+ if block_state.output_type == "latent":
+ block_state.audio = audio_latents
+ else:
+ audio = components.audio_vae.decode(audio_latents, return_dict=False)[0]
+ block_state.audio = audio.float().permute(1, 0, 2)
+ block_state.sampling_rate = components.audio_sampling_rate
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/simpletuner/helpers/models/minimaxh3/denoise.py b/simpletuner/helpers/models/minimaxh3/denoise.py
new file mode 100644
index 000000000..7e8a2887c
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/denoise.py
@@ -0,0 +1,571 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from diffusers.modular_pipelines.modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.utils import logging
+
+from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline
+from .scheduler import MiniMaxH3Scheduler
+from .transformer import MiniMaxH3Transformer3DModel
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def optimized_scale(positive_flat: torch.Tensor, negative_flat: torch.Tensor) -> torch.Tensor:
+ dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+ squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+ return dot_product / squared_norm
+
+
+def _denoiser_inputs() -> list[InputParam]:
+ r"""Everything one MiniMax-H3 forward reads, beyond the transformer itself."""
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The video rows of the packed sequence, conditioning rows first.",
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The channel-major audio rows of the packed sequence, reference rows first.",
+ ),
+ InputParam.template("prompt_embeds"),
+ InputParam(
+ name="negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Optional negative prompt embeddings for real CFG.",
+ ),
+ InputParam(
+ name="row_timestep_plan",
+ type_hint=list,
+ required=True,
+ description="One `(timestep, timestep_indices)` pair per step.",
+ ),
+ InputParam(name="token_tags", type_hint=torch.Tensor, required=True, description="The modality tag of every row."),
+ InputParam(
+ name="position_ids",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The `(t, h, w)` rotary coordinate of every row.",
+ ),
+ InputParam(
+ name="video_indices",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Sequence positions of the video rows.",
+ ),
+ InputParam(
+ name="audio_indices",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Sequence positions of the audio rows.",
+ ),
+ InputParam(
+ name="text_indices",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Sequence positions of the text rows.",
+ ),
+ InputParam(
+ name="negative_row_timestep_plan",
+ type_hint=list,
+ description="Optional negative-branch row-to-timestep plan for real CFG.",
+ ),
+ InputParam(
+ name="negative_token_tags",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch modality tag of every row.",
+ ),
+ InputParam(
+ name="negative_position_ids",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch `(t, h, w)` rotary coordinate of every row.",
+ ),
+ InputParam(
+ name="negative_video_indices",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch sequence positions of the video rows.",
+ ),
+ InputParam(
+ name="negative_audio_indices",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch sequence positions of the audio rows.",
+ ),
+ InputParam(
+ name="negative_text_indices",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch sequence positions of the text rows.",
+ ),
+ InputParam(name="guidance_scale", type_hint=float, default=1.0, description="Real CFG scale."),
+ InputParam(
+ name="guidance_scale_real",
+ type_hint=float,
+ description="Alias for `guidance_scale`, matching SimpleTuner validation configuration.",
+ ),
+ InputParam(
+ name="guidance_rescale",
+ type_hint=float,
+ description="Optional standard-deviation guidance rescale factor.",
+ ),
+ InputParam(
+ name="skip_guidance_layers",
+ type_hint=list,
+ description="Layer indices to skip when computing skipped-layer guidance.",
+ ),
+ InputParam(
+ name="skip_layer_guidance_scale",
+ type_hint=float,
+ default=2.8,
+ description="Scale for skipped-layer guidance.",
+ ),
+ InputParam(
+ name="skip_layer_guidance_start",
+ type_hint=float,
+ default=0.01,
+ description="Fraction of denoising steps after which skipped-layer guidance starts.",
+ ),
+ InputParam(
+ name="skip_layer_guidance_stop",
+ type_hint=float,
+ default=0.2,
+ description="Fraction of denoising steps before which skipped-layer guidance stops.",
+ ),
+ InputParam(
+ name="use_cfg_zero_star",
+ type_hint=bool,
+ default=True,
+ description="Whether to apply CFG Zero* when real CFG is active.",
+ ),
+ InputParam(
+ name="use_zero_init",
+ type_hint=bool,
+ default=True,
+ description="Whether CFG Zero* should zero the first guided steps.",
+ ),
+ InputParam(
+ name="zero_steps",
+ type_hint=int,
+ default=0,
+ description="Number of initial guided steps to zero when `use_zero_init` is enabled.",
+ ),
+ InputParam(
+ name="no_cfg_until_timestep",
+ type_hint=int,
+ default=0,
+ description="Step index before which real CFG is disabled.",
+ ),
+ InputParam(
+ name="cfg_end_timestep",
+ type_hint=int,
+ description="Last step index that may use real CFG.",
+ ),
+ InputParam(
+ name="minimax_h3_reference_mode",
+ type_hint=str,
+ default="vanilla",
+ description="MiniMax-H3 static reference handling mode: vanilla or cached_kv.",
+ ),
+ InputParam.template("attention_kwargs"),
+ ]
+
+
+def _denoiser_outputs() -> list[OutputParam]:
+ return [
+ OutputParam(
+ "noise_pred", type_hint=torch.Tensor, description="Predicted velocity of the video rows of the sequence."
+ ),
+ OutputParam(
+ "audio_noise_pred",
+ type_hint=torch.Tensor,
+ description="Predicted velocity of the audio rows of the sequence.",
+ ),
+ ]
+
+
+def _state_attr(block_state: BlockState, name: str, prefix: str = ""):
+ value = getattr(block_state, f"{prefix}{name}", None) if prefix else getattr(block_state, name)
+ if value is None and prefix:
+ return getattr(block_state, name)
+ return value
+
+
+def _state_attr_or(block_state: BlockState, name: str, default, prefix: str = ""):
+ value = getattr(block_state, f"{prefix}{name}", None) if prefix else getattr(block_state, name, None)
+ if value is None and prefix:
+ value = getattr(block_state, name, None)
+ return default if value is None else value
+
+
+def _predict_velocity(
+ transformer: MiniMaxH3Transformer3DModel,
+ block_state: BlockState,
+ i: int,
+ *,
+ prefix: str = "",
+ prompt_embeds: torch.Tensor | None = None,
+ skip_layers: list[int] | None = None,
+):
+ r"""One MiniMax-H3 forward pass: every row of the packed sequence, at its own noise level, at once."""
+ row_timestep_plan = _state_attr(block_state, "row_timestep_plan", prefix)
+ unique_timesteps, timestep_indices = row_timestep_plan[i]
+ prompt_embeds = block_state.prompt_embeds if prompt_embeds is None else prompt_embeds
+ return transformer(
+ hidden_states=block_state.latents[None],
+ audio_hidden_states=block_state.audio_latents[None],
+ encoder_hidden_states=prompt_embeds,
+ timestep=unique_timesteps,
+ timestep_indices=timestep_indices,
+ token_tags=_state_attr(block_state, "token_tags", prefix),
+ position_ids=_state_attr(block_state, "position_ids", prefix),
+ video_indices=_state_attr(block_state, "video_indices", prefix),
+ audio_indices=_state_attr(block_state, "audio_indices", prefix),
+ text_indices=_state_attr(block_state, "text_indices", prefix),
+ attention_kwargs=getattr(block_state, "attention_kwargs", None),
+ skip_layers=skip_layers,
+ num_condition_video_rows=_state_attr_or(block_state, "num_condition_video_rows", 0, prefix),
+ num_condition_audio_rows=_state_attr_or(block_state, "num_condition_audio_rows", 0, prefix),
+ minimax_h3_reference_mode=getattr(block_state, "minimax_h3_reference_mode", "vanilla") or "vanilla",
+ return_dict=False,
+ )
+
+
+def _resolve_guidance_scale(block_state: BlockState) -> float:
+ guidance_scale_real = getattr(block_state, "guidance_scale_real", None)
+ if guidance_scale_real is not None:
+ return float(guidance_scale_real)
+ return float(getattr(block_state, "guidance_scale", 1.0) or 1.0)
+
+
+def _within_cfg_window(block_state: BlockState, i: int) -> bool:
+ start = int(getattr(block_state, "no_cfg_until_timestep", 0) or 0)
+ stop = getattr(block_state, "cfg_end_timestep", None)
+ return i >= start and (stop is None or i <= int(stop))
+
+
+def _within_skip_window(block_state: BlockState, i: int, num_steps: int) -> bool:
+ start = float(getattr(block_state, "skip_layer_guidance_start", 0.01) or 0.0)
+ stop = float(getattr(block_state, "skip_layer_guidance_stop", 0.2) or 0.0)
+ return i > num_steps * start and i < num_steps * stop
+
+
+def _apply_cfg(
+ positive: torch.Tensor,
+ negative: torch.Tensor,
+ scale: float,
+ block_state: BlockState,
+ i: int,
+) -> torch.Tensor:
+ if bool(getattr(block_state, "use_cfg_zero_star", True)):
+ positive_flat = positive.reshape(positive.shape[0], -1).float()
+ negative_flat = negative.reshape(negative.shape[0], -1).float()
+ alpha = optimized_scale(positive_flat, negative_flat).view(positive.shape[0], *([1] * (positive.ndim - 1)))
+ alpha = alpha.to(device=positive.device, dtype=positive.dtype)
+ if i <= int(getattr(block_state, "zero_steps", 0) or 0) and bool(getattr(block_state, "use_zero_init", True)):
+ return positive * 0.0
+ return negative * alpha + scale * (positive - negative * alpha)
+ return positive + (scale - 1.0) * (positive - negative)
+
+
+def _rescale_guidance(guided: torch.Tensor, positive: torch.Tensor, guidance_rescale: float | None) -> torch.Tensor:
+ if guidance_rescale is None or float(guidance_rescale) <= 0.0:
+ return guided
+ axes = tuple(range(1, guided.ndim))
+ std_positive = torch.std(positive.float(), dim=axes, keepdim=True)
+ std_guided = torch.std(guided.float(), dim=axes, keepdim=True)
+ factor = (std_positive / (std_guided + 1e-8)).to(device=guided.device, dtype=guided.dtype)
+ rescale = float(guidance_rescale)
+ return guided * (1.0 - rescale + rescale * factor)
+
+
+def _validate_negative_branch(block_state: BlockState) -> torch.Tensor:
+ negative_prompt_embeds = getattr(block_state, "negative_prompt_embeds", None)
+ if negative_prompt_embeds is None:
+ raise ValueError("MiniMax-H3 real CFG requires `negative_prompt` or `negative_prompt_embeds`.")
+ negative_text_indices = _state_attr(block_state, "text_indices", "negative_")
+ expected_text_rows = int(negative_text_indices.shape[0])
+ if negative_prompt_embeds.shape[1] != expected_text_rows:
+ raise ValueError(
+ "MiniMax-H3 negative prompt embeds must match their packed text layout: "
+ f"got {negative_prompt_embeds.shape[1]} embeds for {expected_text_rows} text rows."
+ )
+ return negative_prompt_embeds
+
+
+def _predict_guided_velocity(
+ transformer: MiniMaxH3Transformer3DModel,
+ block_state: BlockState,
+ i: int,
+ num_steps: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ positive_video, positive_audio = _predict_velocity(transformer, block_state, i)
+ video_velocity, audio_velocity = positive_video, positive_audio
+
+ guidance_scale = _resolve_guidance_scale(block_state)
+ apply_cfg = guidance_scale > 1.0 and _within_cfg_window(block_state, i)
+ if apply_cfg:
+ negative_prompt_embeds = _validate_negative_branch(block_state)
+ negative_video, negative_audio = _predict_velocity(
+ transformer,
+ block_state,
+ i,
+ prefix="negative_",
+ prompt_embeds=negative_prompt_embeds,
+ )
+ video_velocity = _apply_cfg(positive_video, negative_video, guidance_scale, block_state, i)
+ audio_velocity = _apply_cfg(positive_audio, negative_audio, guidance_scale, block_state, i)
+ guidance_rescale = getattr(block_state, "guidance_rescale", None)
+ video_velocity = _rescale_guidance(video_velocity, positive_video, guidance_rescale)
+ audio_velocity = _rescale_guidance(audio_velocity, positive_audio, guidance_rescale)
+
+ skip_guidance_layers = getattr(block_state, "skip_guidance_layers", None)
+ if skip_guidance_layers is not None:
+ skip_guidance_layers = list(skip_guidance_layers)
+ if skip_guidance_layers and _within_skip_window(block_state, i, num_steps):
+ skip_video, skip_audio = _predict_velocity(
+ transformer,
+ block_state,
+ i,
+ skip_layers=skip_guidance_layers,
+ )
+ scale = float(getattr(block_state, "skip_layer_guidance_scale", 2.8) or 0.0)
+ video_velocity = video_velocity + (positive_video - skip_video) * scale
+ audio_velocity = audio_velocity + (positive_audio - skip_audio) * scale
+
+ return video_velocity, audio_velocity
+
+
+class MiniMaxH3LoopDenoiser(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Runs the one MiniMax-H3 forward pass of a denoising iteration, which predicts the velocity of every row "
+ "of the packed sequence at once. The checkpoint is guidance-distilled, so there is no unconditional pass "
+ "and no guider."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", MiniMaxH3Transformer3DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return _denoiser_inputs()
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return _denoiser_outputs()
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.noise_pred, block_state.audio_noise_pred = _predict_guided_velocity(
+ components.transformer, block_state, i, len(block_state.timesteps)
+ )
+ return components, block_state
+
+
+class MiniMaxH3Ref2VALoopDenoiser(ModularPipelineBlocks):
+ model_name = "minimax-h3-ref2va"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Runs the one MiniMax-H3 forward pass of a `ref2va` denoising iteration, against the `transformer_ref` "
+ "partition of the checkpoint."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer_ref", MiniMaxH3Transformer3DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return _denoiser_inputs()
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return _denoiser_outputs()
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.noise_pred, block_state.audio_noise_pred = _predict_guided_velocity(
+ components.transformer_ref, block_state, i, len(block_state.timesteps)
+ )
+ return components, block_state
+
+
+class MiniMaxH3LoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Steps the generated video and audio rows down their own schedule. The conditioning rows are re-imposed "
+ "by construction: only the generated rows are ever written, so the anchors survive the whole loop."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The video rows of the packed sequence, conditioning rows first.",
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The channel-major audio rows of the packed sequence, reference rows first.",
+ ),
+ InputParam(
+ name="noise_pred",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Predicted velocity of the video rows.",
+ ),
+ InputParam(
+ name="audio_noise_pred",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Predicted velocity of the audio rows.",
+ ),
+ InputParam(
+ name="audio_timesteps",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Timesteps of the audio schedule.",
+ ),
+ InputParam(
+ name="num_condition_video_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading video rows are conditioning rows.",
+ ),
+ InputParam(
+ name="num_condition_audio_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading audio rows are reference rows.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The video rows of the packed sequence after one step.",
+ ),
+ OutputParam(
+ "audio_latents",
+ type_hint=torch.Tensor,
+ description="The audio rows of the packed sequence after one step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ num_condition_video_rows = block_state.num_condition_video_rows
+ num_condition_audio_rows = block_state.num_condition_audio_rows
+
+ block_state.latents[num_condition_video_rows:] = components.scheduler.step(
+ block_state.noise_pred[0, num_condition_video_rows:].float(),
+ t,
+ block_state.latents[num_condition_video_rows:],
+ return_dict=False,
+ )[0]
+ has_audio_rows = block_state.audio_latents is not None and int(block_state.audio_latents.shape[0]) > int(
+ num_condition_audio_rows
+ )
+ if has_audio_rows:
+ if components.audio_scheduler is None:
+ raise ValueError("MiniMax-H3 audio rows require an audio scheduler.")
+ block_state.audio_latents[num_condition_audio_rows:] = components.audio_scheduler.step(
+ block_state.audio_noise_pred[0, num_condition_audio_rows:].float(),
+ block_state.audio_timesteps[i],
+ block_state.audio_latents[num_condition_audio_rows:],
+ return_dict=False,
+ )[0]
+ return components, block_state
+
+
+class MiniMaxH3DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return "Iteratively denoises the packed MiniMax-H3 sequence over the two schedules."
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("timesteps", required=True, description="Timesteps of the video schedule."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ with self.progress_bar(total=len(block_state.timesteps)) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ progress_bar.update()
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3DenoiseStep(MiniMaxH3DenoiseLoopWrapper):
+ block_classes = [MiniMaxH3LoopDenoiser, MiniMaxH3LoopSchedulerStep]
+ block_names = ["denoiser", "update"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the `t2va` / `fl2va` MiniMax-H3 denoising loop, one forward pass per step."
+
+
+class MiniMaxH3Ref2VADenoiseStep(MiniMaxH3DenoiseLoopWrapper):
+ model_name = "minimax-h3-ref2va"
+ block_classes = [MiniMaxH3Ref2VALoopDenoiser, MiniMaxH3LoopSchedulerStep]
+ block_names = ["denoiser", "update"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the `ref2va` MiniMax-H3 denoising loop, one forward pass per step."
diff --git a/simpletuner/helpers/models/minimaxh3/encoders.py b/simpletuner/helpers/models/minimaxh3/encoders.py
new file mode 100644
index 000000000..bc2b1908a
--- /dev/null
+++ b/simpletuner/helpers/models/minimaxh3/encoders.py
@@ -0,0 +1,865 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.utils import logging
+from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+
+from .autoencoder import AutoencoderKLMiniMaxH3
+from .autoencoder_audio import AutoencoderKLMiniMaxH3Audio
+from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline
+from .packing import (
+ MINIMAX_H3_KEYFRAME_NOISE_AUG,
+ MINIMAX_H3_PIXEL_MEAN,
+ MINIMAX_H3_PIXEL_STD,
+ MINIMAX_H3_TEXT_ENCODER_LAYER,
+ MINIMAX_H3_TEXT_TAG,
+ MINIMAX_H3_VIDEO_TAG,
+ keyframe_condition_noise,
+ patchify_video_latents,
+)
+from .packing_ref2va import (
+ MiniMaxH3PreparedReference,
+ build_ref2va_presentation,
+ sample_reference_video_frames,
+ trim_reference_num_frames,
+)
+from .scheduler import MiniMaxH3Scheduler
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH = 512
+
+
+def _check_prompt(prompt) -> None:
+ r"""MiniMax-H3 packs one request into one sequence, so a batch of prompts is not a thing."""
+ if not isinstance(prompt, str):
+ raise ValueError(
+ f"MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got {type(prompt)}."
+ )
+
+
+def _null_token_id(tokenizer) -> int:
+ null_token_id = getattr(tokenizer, "pad_token_id", None)
+ if null_token_id is None:
+ null_token_id = getattr(tokenizer, "eos_token_id", None)
+ if null_token_id is None:
+ raise ValueError("MiniMax-H3 null conditioning requires a tokenizer pad token or eos token.")
+ return int(null_token_id)
+
+
+def _resolve_max_text_length(max_length: int | None) -> int | None:
+ if max_length is None:
+ return MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH
+ max_length = int(max_length)
+ if max_length <= 0:
+ return None
+ return max_length
+
+
+def _encode_instruction_token_ids(tokenizer, prompt: str, null_instruction: bool, max_length: int | None) -> list[int]:
+ prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
+ max_length = _resolve_max_text_length(max_length)
+ if max_length is not None:
+ prompt_ids = prompt_ids[:max_length]
+ if null_instruction:
+ prompt_ids = [_null_token_id(tokenizer)] * len(prompt_ids)
+ return prompt_ids
+
+
+def _conditioner_components() -> list[ComponentSpec]:
+ r"""MiniMax-H3's conditioner: a Qwen3-VL read at its 50th decoder layer, with its language-model head unused."""
+ return [
+ ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration),
+ ComponentSpec("tokenizer", Qwen2TokenizerFast),
+ ComponentSpec("processor", Qwen3VLProcessor),
+ ]
+
+
+def _conditioner_outputs() -> list[OutputParam]:
+ return [
+ OutputParam.template(
+ "prompt_embeds",
+ description=(
+ "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the "
+ "50th decoder layer of the Qwen3-VL conditioner."
+ ),
+ ),
+ OutputParam(
+ "text_token_tags",
+ type_hint=torch.Tensor,
+ description="The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video.",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Optional negative-branch hidden states for real CFG.",
+ ),
+ OutputParam(
+ "negative_text_token_tags",
+ type_hint=torch.Tensor,
+ description="Optional per-row modality tags for `negative_prompt_embeds`.",
+ ),
+ ]
+
+
+class MiniMaxH3TextEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes MiniMax-H3's presentation of a `t2va` / `fl2va` request: the prompt verbatim, preceded by a "
+ '`": "` label and a vision block per keyframe, with no chat template and no special tokens. '
+ "The checkpoint is guidance-distilled, so there is no negative prompt and no unconditional branch."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return _conditioner_components()
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="prompt",
+ type_hint=str,
+ default=None,
+ description="The prompt to guide generation, a single string.",
+ ),
+ InputParam(
+ name="prompt_embeds",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Precomputed MiniMax-H3 text conditioning.",
+ ),
+ InputParam(
+ name="text_token_tags",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Per-row modality tags for precomputed `prompt_embeds`.",
+ ),
+ InputParam(
+ name="negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Optional precomputed negative branch for real CFG.",
+ ),
+ InputParam(
+ name="negative_text_token_tags",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Per-row modality tags for precomputed `negative_prompt_embeds`.",
+ ),
+ InputParam(
+ name="negative_prompt",
+ type_hint=str,
+ description="Optional negative prompt used only when real CFG is enabled.",
+ ),
+ InputParam(
+ name="keyframes",
+ type_hint=list,
+ description="The keyframes put onto the target canvas, in packed order (empty or None for `t2va`).",
+ ),
+ InputParam(
+ name="max_sequence_length",
+ type_hint=int,
+ default=MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH,
+ description="Maximum caption tokens to encode. Vision blocks are structural and are never truncated.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return _conditioner_outputs()
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt: str,
+ images: list | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ null_instruction: bool = False,
+ max_length: int | None = MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ r"""
+ Build MiniMax-H3's presentation of a request and encode it.
+
+ The presentation is the verbatim prompt for `t2va`. Every keyframe prepends a `": "` label and a
+ vision block (`<|vision_start|>`, one `<|image_pad|>` per vision patch, `<|vision_end|>`) — no chat template
+ and no special tokens. The rows of a vision block are tagged as *video* rather than text, which is what the
+ transformer's AdaLN modulation keys off.
+
+ Args:
+ prompt (`str`): The prompt to encode.
+ images (`list[PIL.Image.Image]`, *optional*):
+ The keyframes, already prepared onto the target canvas, in packed order.
+ device (`torch.device`, *optional*): The device to run the conditioner on.
+ dtype (`torch.dtype`, *optional*): The dtype of the returned embeddings.
+ null_instruction (`bool`, *optional*):
+ Replace only the final prompt's token ids with tokenizer null tokens while keeping the prompt token
+ count unchanged. This keeps media rotary positions aligned for H3 real-CFG null branches.
+ max_length (`int`, *optional*, defaults to 512):
+ Maximum prompt/caption tokens to encode. Set `0` or a negative value to disable this cap. Keyframe
+ labels and vision blocks are structural conditioning and are never truncated.
+
+ Returns:
+ `tuple[torch.Tensor, torch.Tensor]`: the `(1, num_text_tokens, 5120)` hidden states and the
+ `(num_text_tokens,)` per-row modality tags.
+ """
+ return MiniMaxH3TextEncoderStep.encode_prompt_batch(
+ components,
+ [prompt],
+ image_batches=[images],
+ device=device,
+ dtype=dtype,
+ null_instructions=[null_instruction],
+ max_length=max_length,
+ )[0]
+
+ @staticmethod
+ def encode_prompt_batch(
+ components,
+ prompts: list[str],
+ image_batches: list[list | None] | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ null_instructions: list[bool] | None = None,
+ max_length: int | None = MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH,
+ ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+ """Encode multiple independently packed H3 requests in one padded Qwen3-VL forward."""
+ if not prompts:
+ return []
+ for prompt in prompts:
+ _check_prompt(prompt)
+ if image_batches is None:
+ image_batches = [None] * len(prompts)
+ if null_instructions is None:
+ null_instructions = [False] * len(prompts)
+ if len(image_batches) != len(prompts) or len(null_instructions) != len(prompts):
+ raise ValueError("MiniMax-H3 batched text encoding requires one image and null-instruction entry per prompt.")
+
+ device = device or components._execution_device
+ dtype = dtype or components.transformer.dtype
+
+ num_layers = components.text_encoder.config.text_config.num_hidden_layers
+ if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER:
+ raise ValueError(
+ f"MiniMax-H3 conditions on `hidden_states[{MINIMAX_H3_TEXT_ENCODER_LAYER}]` of its Qwen3-VL "
+ f"conditioner, which needs more than {MINIMAX_H3_TEXT_ENCODER_LAYER} decoder layers, but "
+ f"`text_encoder` has {num_layers}. The last hidden state of a stack truncated to exactly "
+ f"{MINIMAX_H3_TEXT_ENCODER_LAYER} layers is post-norm and is not the conditioning MiniMax-H3 expects."
+ )
+
+ presentations = []
+ flattened_images = [image for images in image_batches if images for image in images]
+ if flattened_images:
+ vision = components.processor.image_processor(images=flattened_images, return_tensors="pt")
+ pixel_values = vision["pixel_values"]
+ image_grid_thw = vision["image_grid_thw"]
+ else:
+ pixel_values, image_grid_thw = None, None
+ image_grid_offset = 0
+ for prompt, images, null_instruction in zip(prompts, image_batches, null_instructions):
+ token_ids, token_tags = [], []
+ if images:
+ request_image_grid_thw = image_grid_thw[image_grid_offset : image_grid_offset + len(images)]
+ image_grid_offset += len(images)
+ merge_size = components.processor.image_processor.merge_size**2
+ for index in range(len(images)):
+ num_image_tokens = int(request_image_grid_thw[index].prod()) // merge_size
+ label_ids = components.tokenizer(f": ", add_special_tokens=False)["input_ids"]
+ vision_ids = (
+ [components.tokenizer.convert_tokens_to_ids("<|vision_start|>")]
+ + [components.tokenizer.convert_tokens_to_ids("<|image_pad|>")] * num_image_tokens
+ + [components.tokenizer.convert_tokens_to_ids("<|vision_end|>")]
+ )
+ token_ids += label_ids + vision_ids
+ token_tags += [MINIMAX_H3_TEXT_TAG] * len(label_ids) + [MINIMAX_H3_VIDEO_TAG] * len(vision_ids)
+ prompt_ids = _encode_instruction_token_ids(components.tokenizer, prompt, null_instruction, max_length)
+ token_ids += prompt_ids
+ token_tags += [MINIMAX_H3_TEXT_TAG] * len(prompt_ids)
+ if not token_ids:
+ raise ValueError("MiniMax-H3 conditioning carries no tokens; conditioning cannot be empty.")
+ presentations.append((token_ids, token_tags))
+
+ max_sequence_length = max(len(token_ids) for token_ids, _ in presentations)
+ input_ids = torch.full(
+ (len(prompts), max_sequence_length),
+ _null_token_id(components.tokenizer),
+ dtype=torch.long,
+ device=device,
+ )
+ attention_mask = torch.zeros_like(input_ids)
+ mm_token_type_ids = torch.zeros_like(input_ids)
+ for batch_index, (token_ids, _) in enumerate(presentations):
+ sequence_length = len(token_ids)
+ input_ids[batch_index, :sequence_length] = torch.tensor(token_ids, dtype=torch.long, device=device)
+ attention_mask[batch_index, :sequence_length] = 1
+ # Qwen3-VL reads these token types to construct per-sample multimodal rotary positions.
+ token_type_ids = components.processor.create_mm_token_type_ids([token_ids])[0]
+ mm_token_type_ids[batch_index, :sequence_length] = torch.tensor(token_type_ids, dtype=torch.long, device=device)
+
+ # `text_encoder.model` is a submodule, and a CPU-offload hook — accelerate's or the one the
+ # `ComponentsManager` attaches — wraps the *top-level* module's `forward` alone, so calling the submodule
+ # directly would leave the conditioner on the CPU. Fire the hook by hand instead of routing through
+ # `text_encoder(...)`: MiniMax-H3 reads `hidden_states[50]` and never uses the language-model head, whose
+ # vocabulary-wide projection over every token is all the top-level forward would add.
+ hook = getattr(components.text_encoder, "_hf_hook", None)
+ if hook is not None and hasattr(hook, "pre_forward"):
+ hook.pre_forward(components.text_encoder)
+ outputs = components.text_encoder.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ mm_token_type_ids=mm_token_type_ids,
+ pixel_values=None if pixel_values is None else pixel_values.to(device, components.text_encoder.dtype),
+ image_grid_thw=None if image_grid_thw is None else image_grid_thw.to(device),
+ use_cache=False,
+ output_hidden_states=True,
+ )
+ hidden_states = outputs.hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER]
+ return [
+ (
+ hidden_states[batch_index : batch_index + 1, : len(token_ids)].to(device=device, dtype=dtype),
+ torch.tensor(token_tags, dtype=torch.long),
+ )
+ for batch_index, (token_ids, token_tags) in enumerate(presentations)
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ if getattr(block_state, "prompt_embeds", None) is not None:
+ if getattr(block_state, "text_token_tags", None) is None:
+ raise ValueError("MiniMax-H3 precomputed `prompt_embeds` require `text_token_tags`.")
+ self.set_block_state(state, block_state)
+ return components, state
+
+ _check_prompt(block_state.prompt)
+
+ # `encode_prompt` defaults the embedding dtype to the denoiser's; a text encoder block has no denoiser of
+ # its own — it is meant to run on its own — so it emits the conditioner's dtype, as every other model does.
+ block_state.prompt_embeds, block_state.text_token_tags = self.encode_prompt(
+ components,
+ block_state.prompt,
+ block_state.keyframes,
+ device=components._execution_device,
+ dtype=components.text_encoder.dtype,
+ max_length=getattr(block_state, "max_sequence_length", MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH),
+ )
+ if getattr(block_state, "negative_prompt", None) is not None:
+ _check_prompt(block_state.negative_prompt)
+ block_state.negative_prompt_embeds, block_state.negative_text_token_tags = self.encode_prompt(
+ components,
+ block_state.prompt if block_state.negative_prompt.strip() == "" else block_state.negative_prompt,
+ block_state.keyframes,
+ device=components._execution_device,
+ dtype=components.text_encoder.dtype,
+ null_instruction=block_state.negative_prompt.strip() == "",
+ max_length=getattr(block_state, "max_sequence_length", MINIMAX_H3_DEFAULT_MAX_TEXT_LENGTH),
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3KeyframeVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes the `fl2va` keyframes into packed conditioning rows and noises them to MiniMax-H3's "
+ "conditioning level. The rows are the anchors of the whole denoising loop: the loop only ever writes the "
+ "generated rows, so they are never updated again."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLMiniMaxH3),
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="keyframes",
+ type_hint=list,
+ required=True,
+ description="The keyframes put onto the target canvas, in packed order.",
+ ),
+ InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."),
+ InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."),
+ InputParam.template(
+ "generator",
+ description=(
+ "The generator of the request. The conditioning noise is drawn from it before the target noise "
+ "of the prepare-latents step."
+ ),
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "condition_latents",
+ type_hint=torch.Tensor,
+ description="The noise-augmented video conditioning rows, in packed order.",
+ )
+ ]
+
+ @staticmethod
+ def encode_keyframes(components, images: list, device: torch.device | None = None) -> torch.Tensor:
+ r"""
+ Encode the `fl2va` keyframes into packed conditioning rows.
+
+ The keyframes go through the video VAE's spatial encoder only — they are single frames, so none of its
+ 17-frame temporal chunking applies — and the posterior mean is used, matching ComfyUI's MiniMax-H3 VAE path.
+
+ Args:
+ images (`list[PIL.Image.Image]`):
+ The keyframes, already prepared onto the target canvas, in packed order.
+ device (`torch.device`, *optional*): The device to run the VAE on.
+
+ Returns:
+ `torch.Tensor` of shape `(num_condition_rows, latent_channels * prod(patch_size))`: the float32
+ conditioning rows.
+ """
+ device = device or components._execution_device
+ latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1)
+ latents_std = torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1)
+ pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1)
+ pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1)
+
+ rows = []
+ for image in images:
+ pixels = torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None]
+ pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std
+ # `vae.encode` chunks along time for videos; a keyframe is one frame and is encoded by the (tiled)
+ # spatial encoder alone, which is what the released model conditions on.
+ moments = components.vae._encode_clip(pixels)
+ posterior = DiagonalGaussianDistribution(moments)
+ latents = posterior.mode().float().cpu()
+ rows.append(patchify_video_latents((latents - latents_mean) / latents_std, components.patch_size))
+ return torch.cat(rows)
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ condition_latents = self.encode_keyframes(components, block_state.keyframes, device=device)
+ noise = keyframe_condition_noise(
+ ((1, block_state.latent_height, block_state.latent_width),) * len(block_state.keyframes),
+ components.patch_size,
+ components.vae_latent_channels,
+ generator=block_state.generator,
+ device=device,
+ )
+ block_state.condition_latents = components.scheduler.scale_noise(
+ condition_latents.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VATextEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3-ref2va"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes MiniMax-H3's presentation of a `ref2va` request: a label per reference, numbered per modality "
+ '(`": "` plus a vision block, `"