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
                 )