github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-05 13:27:12 -06:00
parent 9065951da3
commit 1e1418b22c
8 changed files with 163 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -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(

122
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
View File

@@ -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
)

20
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -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()

5
toolkit/config_modules.py
View File

@@ -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)

4
toolkit/network_mixins.py
View File

@@ -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)

4
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -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}

4
ui/src/app/jobs/new/options.tsx
View File

@@ -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>

2
version.py
View File

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