Apply sigma to contrastive guidance to balance loss better. Prevent noise grads on images/non audio datasets. Prep for training adapters on MiniMax H3

extensions_built_in/diffusion_models/ltx2/ltx2.py

@@ -942,7 +942,17 @@ class LTX2Model(BaseModel):
                 patch_size_t=self.pipeline.transformer_temporal_patch_size,
             )
 
-            if batch.audio_latents is not None or batch.audio_tensor is not None:
+            # audio only trains for video batches from datasets that asked for
+            # it. Cached latents can carry audio after do_audio was turned off,
+            # and image (single frame) batches must never pick up a soundtrack.
+            do_audio = (
+                batch.dataset_config is not None
+                and batch.dataset_config.do_audio
+                and getattr(batch, "num_frames", 1) > 1
+            )
+            if do_audio and (
+                batch.audio_latents is not None or batch.audio_tensor is not None
+            ):
                 if batch.audio_latents is not None:
                     # we have audio latents cached
                     raw_audio_latents = batch.audio_latents.to(

extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py

@@ -276,6 +276,110 @@ class MinimaxH3Model(BaseModel):
             )
         return f"dit_{partition}"
 
+    def load_training_adapter(self, transformer: MiniMaxH3Transformer):
+        """Load an assistant LoRA (e.g. a de-distillation adapter) as a LIVE
+        module: active during training, deactivated by the sampler. It is
+        deliberately NOT merged into the base weights — the transformer is
+        pre-quantized, and a merge would resample every int8 scale.
+
+        Path resolution: a local path is used as-is; otherwise the loras
+        folder under MODELS_PATH is searched recursively for the filename;
+        otherwise a ``user/repo/file.safetensors`` hub path downloads into
+        MODELS_PATH/loras/training_adapters/.
+        """
+        from toolkit.config_modules import NetworkConfig
+        from toolkit.lora_special import LoRASpecialNetwork
+
+        self.print_and_status_update("Loading assistant LoRA")
+        lora_path = self.model_config.assistant_lora_path
+        if not os.path.exists(lora_path):
+            filename = os.path.basename(lora_path)
+            found = self._find_file_recursive(
+                os.path.join(MODELS_PATH, "loras"), filename
+            )
+            if found is not None:
+                lora_path = found
+            else:
+                lora_splits = lora_path.split("/")
+                if len(lora_splits) != 3:
+                    raise ValueError(
+                        f"Assistant LoRA path {lora_path} is not a local path, a "
+                        f"file under {os.path.join(MODELS_PATH, 'loras')}, or a "
+                        "'user/repo/file.safetensors' hub path."
+                    )
+                import huggingface_hub
+
+                target_dir = os.path.join(MODELS_PATH, "loras", "training_adapters")
+                os.makedirs(target_dir, exist_ok=True)
+                try:
+                    lora_path = huggingface_hub.hf_hub_download(
+                        repo_id="/".join(lora_splits[:2]),
+                        filename=lora_splits[2],
+                        local_dir=target_dir,
+                    )
+                except Exception as e:
+                    raise ValueError(
+                        f"Failed to download assistant LoRA from {lora_path}: {e}"
+                    )
+            self.model_config.assistant_lora_path = lora_path
+
+        # load the adapter; it stays a live module (never merged) and the
+        # sampler toggles it off for previews
+        lora_state_dict = load_file(lora_path)
+        dim_key = next(
+            k
+            for k in lora_state_dict
+            if k.endswith("lora_A.weight") or k.endswith("lora_down.weight")
+        )
+        dim = int(lora_state_dict[dim_key].shape[0])
+        lora_state_dict = self.convert_lora_weights_before_load(lora_state_dict)
+
+        network_config = NetworkConfig(
+            **{
+                "type": "lora",
+                "linear": dim,
+                "linear_alpha": dim,
+                "transformer_only": True,
+            }
+        )
+        LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer"
+        network = LoRASpecialNetwork(
+            text_encoder=None,
+            unet=transformer,
+            lora_dim=network_config.linear,
+            multiplier=1.0,
+            alpha=network_config.linear_alpha,
+            train_unet=True,
+            train_text_encoder=False,
+            network_config=network_config,
+            network_type=network_config.type,
+            transformer_only=network_config.transformer_only,
+            is_transformer=True,
+            target_lin_modules=self.target_lora_modules,
+            is_assistant_adapter=True,
+            is_ara=True,
+        )
+        network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True)
+        network.force_to(self.device_torch, dtype=self.torch_dtype)
+        network._update_torch_multiplier()
+        network.load_weights(lora_state_dict)
+
+        # frozen: the adapter shapes the training distribution but is never
+        # itself trained, so its params must not collect gradients
+        network.is_merged_in = False
+        for param in network.parameters():
+            param.requires_grad_(False)
+        network.eval()
+
+        self.assistant_lora: LoRASpecialNetwork = network
+
+        # live during training; the sampler's non-inverted assistant path
+        # (BaseModel.generate_images) deactivates it for previews and turns
+        # it back on afterwards
+        self.assistant_lora.multiplier = 1.0
+        self.assistant_lora.is_active = True
+        self.invert_assistant_lora = False
+
     def _load_transformer(self) -> MiniMaxH3Transformer:
         dtype = self.torch_dtype
         dit_path = self._resolve_comfy_file(self._dit_component())
@@ -460,6 +564,10 @@ class MinimaxH3Model(BaseModel):
 
         transformer = self._load_transformer()
 
+        # load assistant lora if specified (merged into the quantized weights)
+        if self.model_config.assistant_lora_path is not None:
+            self.load_training_adapter(transformer)
+
         if self.model_config.quantize:
             self.print_and_status_update("Quantizing transformer")
             quantize_model(self, transformer)
@@ -779,10 +887,20 @@ class MinimaxH3Model(BaseModel):
                 # invert 17n+5 -> 5n+2 from the latent frame count
                 num_frames = (t_lat - 2) // 5 * 17 + 5 if t_lat > 1 else 1
             a_lat = packing.audio_latent_num_frames(num_frames)
+            # audio only trains for video batches from datasets that asked for
+            # it. Cached latents can carry audio after do_audio was turned off,
+            # and image (single frame) batches must never pick up a soundtrack
+            # — either way it rides along as silence with no audio loss.
+            do_audio = (
+                batch is not None
+                and batch.dataset_config is not None
+                and batch.dataset_config.do_audio
+                and num_frames > 1
+            )
             raw_audio = None
-            if batch is not None and batch.audio_latents is not None:
+            if do_audio and batch.audio_latents is not None:
                 raw_audio = batch.audio_latents.to(device, torch.float32)
-            elif batch is not None and getattr(batch, "audio_data", None) is not None:
+            elif do_audio and getattr(batch, "audio_data", None) is not None:
                 raw_audio = self.encode_audio(batch.audio_data).to(
                     device, torch.float32
                 )

extensions_built_in/sd_trainer/SDTrainer.py

@@ -737,6 +737,14 @@ class SDTrainer(BaseSDTrainProcess):
                     guidance_scale = torch.tensor(guidance_scale).to(target.device, dtype=target.dtype)
                     guidance_scale = guidance_scale.view(-1, 1, 1, 1) if not is_video else guidance_scale.view(-1, 1, 1, 1, 1)
 
+                if self.train_config.guidance_loss_schedule == 'sigma':
+                    # the (target - uncond) sample direction carries s * fresh_noise
+                    # that nothing can predict at low sigma, so decay the
+                    # extrapolation toward a plain flow target as sigma falls
+                    sigma = (timesteps.to(target.device) / 1000.0).to(target.dtype)
+                    sigma = sigma.view(-1, 1, 1, 1) if not is_video else sigma.view(-1, 1, 1, 1, 1)
+                    guidance_scale = 1.0 + (guidance_scale - 1.0) * sigma
+
                 unconditional_target = unconditional_target * alpha
                 target = unconditional_target + guidance_scale * (target - unconditional_target)
 
@@ -762,6 +770,16 @@ class SDTrainer(BaseSDTrainProcess):
                             audio_target.device, dtype=audio_target.dtype
                         ).view(-1, *audio_dims)
 
+                    if self.train_config.guidance_loss_schedule == 'sigma':
+                        # audio streams can run on their own remapped sigma
+                        audio_sigma = getattr(batch, 'audio_sigma', None)
+                        if audio_sigma is None:
+                            audio_sigma = timesteps / 1000.0
+                        audio_sigma = audio_sigma.to(
+                            audio_target.device, dtype=audio_target.dtype
+                        ).view(-1, *audio_dims)
+                        audio_guidance_scale = 1.0 + (audio_guidance_scale - 1.0) * audio_sigma
+
                     batch.audio_target = (
                         audio_uncond + audio_guidance_scale * (audio_target - audio_uncond)
                     ).to(batch.audio_target.dtype).detach()

toolkit/config_modules.py

@@ -590,6 +590,11 @@ class TrainConfig:
         self.do_guidance_loss = kwargs.get('do_guidance_loss', False)
         self.guidance_loss_target: Union[int, List[int, int]] = kwargs.get('guidance_loss_target', 3.0)
         self.do_guidance_loss_cfg_zero: bool = kwargs.get('do_guidance_loss_cfg_zero', False)
+        # 'constant' uses guidance_loss_target as is. 'sigma' decays the target
+        # toward 1.0 as sigma falls (effective = 1 + (target - 1) * sigma) so the
+        # extrapolation never amplifies the unpredictable fresh-noise term at low
+        # sigma. Needed for guidance-distilled models with no guidance embedding.
+        self.guidance_loss_schedule: str = kwargs.get('guidance_loss_schedule', 'sigma')
         self.unconditional_prompt: str = kwargs.get('unconditional_prompt', '')
         if isinstance(self.guidance_loss_target, tuple):
             self.guidance_loss_target = list(self.guidance_loss_target)

toolkit/network_mixins.py

@@ -451,7 +451,9 @@ class ToolkitModuleMixin:
         if is_ao_quantized:
             from toolkit.util.quantize import get_torchao_config, requantize_module_weight
             config = get_torchao_config(self._get_base_qtype())
-            if config is None:
+            if config is None and not getattr(self.org_module[0], "is_ostris_quantized", False):
+                # ostris-quantized layers re-quantize through their own backend
+                # (requantize_module_weight) and need no torchao config
                 print_once(f"Warning: merging into quantized layer {getattr(self, 'lora_name', '?')} "
                            f"without a known qtype; it will be left dequantized")
             requantize_module_weight(self.org_module[0], weight.to(weight_device), orig_dtype, config)

ui/src/app/jobs/new/SimpleJob.tsx

@@ -915,7 +915,7 @@ export default function SimpleJob({
                         if (value) {
                           setJobConfig(true, 'config.process[0].train.do_guidance_loss');
                           if (!jobConfig.config.process[0].train.guidance_loss_target) {
-                            setJobConfig(3.0, 'config.process[0].train.guidance_loss_target');
+                            setJobConfig(4.0, 'config.process[0].train.guidance_loss_target');
                           }
                         } else {
                           setJobConfig(undefined, 'config.process[0].train.do_guidance_loss');
@@ -928,7 +928,7 @@ export default function SimpleJob({
                         <NumberInput
                           label="Guidance Loss Target"
                           docKey={'train.guidance_loss_target'}
-                          value={(jobConfig.config.process[0].train.guidance_loss_target as number) || 3.0}
+                          value={(jobConfig.config.process[0].train.guidance_loss_target as number) || 4.0}
                           onChange={value => setJobConfig(value, 'config.process[0].train.guidance_loss_target')}
                           placeholder="eg. 3.0"
                           min={0}

ui/src/app/jobs/new/options.tsx

@@ -718,7 +718,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
       'config.process[0].train.cache_text_embeddings': [true, false],
       'config.process[0].train.do_guidance_loss': [true, undefined],
-      'config.process[0].train.guidance_loss_target': [3.0, undefined],
+      'config.process[0].train.guidance_loss_target': [4.0, undefined],
       'config.process[0].network.linear': [16, defaultLinearRank],
       'config.process[0].network.linear_alpha': [16, defaultLinearRank],
       'config.process[0].network.network_kwargs.ignore_if_contains': [['adaln_proj'], []],
@@ -738,7 +738,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].datasets[x].auto_frame_count': [true, undefined],
     },
     disableSections: ['network.conv'],
-    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier', 'datasets.auto_frame_count'],
+    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier', 'datasets.auto_frame_count', 'model.assistant_lora_path'],
     modelNotes: (
       <div className="space-y-2">
         <p>

version.py

@@ -1 +1 @@
-VERSION = "0.12.5"
+VERSION = "0.12.6"