diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index c643a0f..bb8ac03 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/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( diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 2ca4d84..0d8367c 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/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 ) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 85829e5..d8adbcc 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -736,7 +736,15 @@ class SDTrainer(BaseSDTrainProcess): if isinstance(guidance_scale, list): 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() diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 961762f..d2b8a59 100644 --- a/toolkit/config_modules.py +++ b/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) diff --git a/toolkit/network_mixins.py b/toolkit/network_mixins.py index 761edee..a23d392 100644 --- a/toolkit/network_mixins.py +++ b/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) diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 5411f78..014b845 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/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({ setJobConfig(value, 'config.process[0].train.guidance_loss_target')} placeholder="eg. 3.0" min={0} diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index 5ef5efb..3be23f7 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/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: (

diff --git a/version.py b/version.py index b6814c0..58fb86d 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.12.5" +VERSION = "0.12.6"