Skip to content

Commit 84f93fc

Browse files
authored
# Add --audio_loss_weight to MiniMax-H3 training (#1655)
* add audio_loss_weight * add reference_max_pixels Size reference videos by their own pixel budget instead of the training canvas: passing --height/--width made ImageCropAndResize center-crop every reference video to the target aspect ratio, which the pipeline then rescaled again to its reference short edge. * document MiniMax-H3 specific training parameters List --silent_on_missing_audio, --reference_max_pixels, --training_cfg_scale and --audio_loss_weight in the zh/en parameter sections, wording taken from the parser help. * revert reference_max_pixels
1 parent 298df41 commit 84f93fc

4 files changed

Lines changed: 16 additions & 4 deletions

File tree

‎diffsynth/diffusion/loss.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ def FlowMatchSFTAudioVideoLoss(pipe: BasePipeline, **inputs):
6363
return loss
6464

6565

66-
def FlowMatchSFTMiniMaxH3AudioVideoLoss(pipe: BasePipeline, training_cfg_scale: float = 1.0, inputs_nega: dict | None = None, **inputs):
66+
def FlowMatchSFTMiniMaxH3AudioVideoLoss(pipe: BasePipeline, training_cfg_scale: float = 1.0, audio_loss_weight: float = 1.0, inputs_nega: dict | None = None, **inputs):
6767
max_timestep_boundary = int(inputs.get("max_timestep_boundary", 1) * len(pipe.scheduler.timesteps))
6868
min_timestep_boundary = int(inputs.get("min_timestep_boundary", 0) * len(pipe.scheduler.timesteps))
6969

@@ -113,7 +113,7 @@ def FlowMatchSFTMiniMaxH3AudioVideoLoss(pipe: BasePipeline, training_cfg_scale:
113113
loss = loss * pipe.scheduler.training_weight(timestep_video)
114114
if "audio_input_latents" in inputs:
115115
loss_audio = torch.nn.functional.mse_loss(noise_pred_audio.float(), training_target_audio.float())
116-
loss_audio = loss_audio * pipe.scheduler_audio.training_weight(timestep_audio)
116+
loss_audio = loss_audio * pipe.scheduler_audio.training_weight(timestep_audio) * audio_loss_weight
117117
loss = loss + loss_audio
118118
return loss
119119

‎docs/en/Model_Details/MiniMax-H3.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,9 @@ Models in the MiniMax-H3 series are trained uniformly via [`examples/minimax_h3/
211211
* MiniMax-H3 Specific Parameters
212212
* `--processor_path`: Path of the Qwen3-VL processor, supports the `model_id:origin_file_pattern` form, used to tokenize the prompt.
213213
* `--initialize_model_on_cpu`: Whether to initialize models on CPU.
214+
* `--silent_on_missing_audio`: Whether to use silent audio as a fallback when no audio track is present in the video data.
215+
* `--training_cfg_scale`: Inverse-CFG scale for preserving MiniMax-H3 guidance distillation during fine-tuning. Values greater than 1 enable a no-grad unconditional branch; 1 keeps the standard flow-matching loss.
216+
* `--audio_loss_weight`: Weight of the audio term in the MiniMax-H3 loss. 1 keeps video and audio equally weighted; 0 trains on the video term only while the audio stream is still noised and forwarded.
214217

215218
We provide an example dataset for testing, which can be downloaded with the following command:
216219

‎docs/zh/Model_Details/MiniMax-H3.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,9 @@ MiniMax-H3 系列模型统一通过 [`examples/minimax_h3/model_training/train.p
211211
* MiniMax-H3 专有参数
212212
* `--processor_path`: Qwen3-VL processor 的路径,支持 `model_id:origin_file_pattern` 形式,用于对 prompt 进行 tokenize。
213213
* `--initialize_model_on_cpu`: 是否在 CPU 上初始化模型。
214+
* `--silent_on_missing_audio`: 视频数据不含音轨时,是否以静音音频作为兜底。
215+
* `--training_cfg_scale`: 微调时用于保留 MiniMax-H3 指导蒸馏的逆 CFG 系数。大于 1 时启用一路无梯度的无条件分支,等于 1 时保持标准 flow matching 损失。
216+
* `--audio_loss_weight`: MiniMax-H3 损失中音频项的权重。为 1 时视频与音频等权重,为 0 时只用视频项训练,音频流仍会加噪并送入模型。
214217

215218
我们构建了一个样例数据集,以方便您进行测试,通过以下命令可以下载这个数据集:
216219

‎examples/minimax_h3/model_training/train.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,15 @@ def __init__(
2929
resume_from_checkpoint=None, remove_prefix_in_ckpt=None,
3030
silent_on_missing_audio=False,
3131
training_cfg_scale=1.0,
32+
audio_loss_weight=1.0,
3233
device="cpu",
3334
task="sft",
3435
):
3536
super().__init__()
3637
if training_cfg_scale < 1.0:
3738
raise ValueError("training_cfg_scale must be at least 1.0")
3839
self.training_cfg_scale = training_cfg_scale
40+
self.audio_loss_weight = audio_loss_weight
3941
# Load models
4042
model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, quant_options=quant_options, device=device)
4143
pipe_kwargs = {}
@@ -69,10 +71,12 @@ def __init__(
6971
self.task_to_loss = {
7072
"sft:data_process": lambda pipe, *args: args,
7173
"sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTMiniMaxH3AudioVideoLoss(
72-
pipe, training_cfg_scale=self.training_cfg_scale, inputs_nega=inputs_nega, **inputs_shared, **inputs_posi,
74+
pipe, training_cfg_scale=self.training_cfg_scale, audio_loss_weight=self.audio_loss_weight,
75+
inputs_nega=inputs_nega, **inputs_shared, **inputs_posi,
7376
),
7477
"sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTMiniMaxH3AudioVideoLoss(
75-
pipe, training_cfg_scale=self.training_cfg_scale, inputs_nega=inputs_nega, **inputs_shared, **inputs_posi,
78+
pipe, training_cfg_scale=self.training_cfg_scale, audio_loss_weight=self.audio_loss_weight,
79+
inputs_nega=inputs_nega, **inputs_shared, **inputs_posi,
7680
),
7781
}
7882

@@ -146,6 +150,7 @@ def minimax_h3_parser():
146150
parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.")
147151
parser.add_argument("--silent_on_missing_audio", default=False, action="store_true", help="Whether to use silent audio as a fallback when no audio track is present in the video data.")
148152
parser.add_argument("--training_cfg_scale", type=float, default=1.0, help="Inverse-CFG scale for preserving MiniMax-H3 guidance distillation during fine-tuning. Values greater than 1 enable a no-grad unconditional branch; 1 keeps the standard flow-matching loss.")
153+
parser.add_argument("--audio_loss_weight", type=float, default=1.0, help="Weight of the audio term in the MiniMax-H3 loss. 1 keeps video and audio equally weighted; 0 trains on the video term only while the audio stream is still noised and forwarded.")
149154
return parser
150155

151156

@@ -220,6 +225,7 @@ def minimax_h3_parser():
220225
remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
221226
silent_on_missing_audio=args.silent_on_missing_audio,
222227
training_cfg_scale=args.training_cfg_scale,
228+
audio_loss_weight=args.audio_loss_weight,
223229
task=args.task,
224230
device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device,
225231
)

0 commit comments

Comments
 (0)