import os import random from collections import OrderedDict from typing import Union, Literal, List, Optional import numpy as np from diffusers import T2IAdapter, AutoencoderTiny, ControlNetModel import torch.functional as F from safetensors.torch import load_file from torch.utils.data import DataLoader, ConcatDataset from toolkit import train_tools from toolkit.basic import value_map, adain, get_mean_std from toolkit.clip_vision_adapter import ClipVisionAdapter from toolkit.config_modules import GenerateImageConfig from toolkit.data_loader import get_dataloader_datasets from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO, FileItemDTO from toolkit.guidance import get_targeted_guidance_loss, get_guidance_loss, GuidanceType from toolkit.image_utils import show_tensors, show_latents from toolkit.ip_adapter import IPAdapter from toolkit.custom_adapter import CustomAdapter from toolkit.print import print_acc from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds from toolkit.reference_adapter import ReferenceAdapter from toolkit.stable_diffusion_model import StableDiffusion, BlankNetwork from toolkit.train_tools import get_torch_dtype, apply_snr_weight, add_all_snr_to_noise_scheduler, \ apply_learnable_snr_gos, LearnableSNRGamma import gc import torch from jobs.process import BaseSDTrainProcess from torchvision import transforms from diffusers import EMAModel import math from toolkit.train_tools import precondition_model_outputs_flow_match from toolkit.models.diffusion_feature_extraction import DiffusionFeatureExtractor, load_dfe from toolkit.util.wavelet_loss import wavelet_loss import torch.nn.functional as F from toolkit.unloader import unload_text_encoder def flush(): torch.cuda.empty_cache() gc.collect() adapter_transforms = transforms.Compose([ transforms.ToTensor(), ]) class SDTrainer(BaseSDTrainProcess): def __init__(self, process_id: int, job, config: OrderedDict, **kwargs): super().__init__(process_id, job, config, **kwargs) self.assistant_adapter: Union['T2IAdapter', 'ControlNetModel', None] self.do_prior_prediction = False self.do_long_prompts = False self.do_guided_loss = False self.taesd: Optional[AutoencoderTiny] = None self._clip_image_embeds_unconditional: Union[List[str], None] = None self.negative_prompt_pool: Union[List[str], None] = None self.batch_negative_prompt: Union[List[str], None] = None self.is_bfloat = self.train_config.dtype == "bfloat16" or self.train_config.dtype == "bf16" self.do_grad_scale = True if self.is_fine_tuning and self.is_bfloat: self.do_grad_scale = False if self.adapter_config is not None: if self.adapter_config.train: self.do_grad_scale = False # if self.train_config.dtype in ["fp16", "float16"]: # # patch the scaler to allow fp16 training # org_unscale_grads = self.scaler._unscale_grads_ # def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16): # return org_unscale_grads(optimizer, inv_scale, found_inf, True) # self.scaler._unscale_grads_ = _unscale_grads_replacer self.cached_blank_embeds: Optional[PromptEmbeds] = None self.cached_trigger_embeds: Optional[PromptEmbeds] = None self.diff_output_preservation_embeds: Optional[PromptEmbeds] = None self.dfe: Optional[DiffusionFeatureExtractor] = None self.unconditional_embeds = None if self.train_config.diff_output_preservation: if self.trigger_word is None: raise ValueError("diff_output_preservation requires a trigger_word to be set") if self.network_config is None: raise ValueError("diff_output_preservation requires a network to be set") if self.train_config.train_text_encoder: raise ValueError("diff_output_preservation is not supported with train_text_encoder") # always do a prior prediction when doing diff output preservation self.do_prior_prediction = True # store the loss target for a batch so we can use it in a loss self._guidance_loss_target_batch: float = 0.0 if isinstance(self.train_config.guidance_loss_target, (int, float)): self._guidance_loss_target_batch = float(self.train_config.guidance_loss_target) elif isinstance(self.train_config.guidance_loss_target, list): self._guidance_loss_target_batch = float(self.train_config.guidance_loss_target[0]) else: raise ValueError(f"Unknown guidance loss target type {type(self.train_config.guidance_loss_target)}") def before_model_load(self): pass def cache_sample_prompts(self): if self.train_config.disable_sampling: return if self.sample_config is not None and self.sample_config.samples is not None and len(self.sample_config.samples) > 0: # cache all the samples self.sd.sample_prompts_cache = [] sample_folder = os.path.join(self.save_root, 'samples') output_path = os.path.join(sample_folder, 'test.jpg') for i in range(len(self.sample_config.prompts)): sample_item = self.sample_config.samples[i] prompt = self.sample_config.prompts[i] # needed so we can autoparse the prompt to handle flags gen_img_config = GenerateImageConfig( prompt=prompt, # it will autoparse the prompt negative_prompt=sample_item.neg, output_path=output_path, ) positive = self.sd.encode_prompt(gen_img_config.prompt).to('cpu') negative = self.sd.encode_prompt(gen_img_config.negative_prompt).to('cpu') self.sd.sample_prompts_cache.append({ 'conditional': positive, 'unconditional': negative }) def before_dataset_load(self): self.assistant_adapter = None # get adapter assistant if one is set if self.train_config.adapter_assist_name_or_path is not None: adapter_path = self.train_config.adapter_assist_name_or_path if self.train_config.adapter_assist_type == "t2i": # dont name this adapter since we are not training it self.assistant_adapter = T2IAdapter.from_pretrained( adapter_path, torch_dtype=get_torch_dtype(self.train_config.dtype) ).to(self.device_torch) elif self.train_config.adapter_assist_type == "control_net": self.assistant_adapter = ControlNetModel.from_pretrained( adapter_path, torch_dtype=get_torch_dtype(self.train_config.dtype) ).to(self.device_torch, dtype=get_torch_dtype(self.train_config.dtype)) else: raise ValueError(f"Unknown adapter assist type {self.train_config.adapter_assist_type}") self.assistant_adapter.eval() self.assistant_adapter.requires_grad_(False) flush() if self.train_config.train_turbo and self.train_config.show_turbo_outputs: if self.model_config.is_xl: self.taesd = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=get_torch_dtype(self.train_config.dtype)) else: self.taesd = AutoencoderTiny.from_pretrained("madebyollin/taesd", torch_dtype=get_torch_dtype(self.train_config.dtype)) self.taesd.to(dtype=get_torch_dtype(self.train_config.dtype), device=self.device_torch) self.taesd.eval() self.taesd.requires_grad_(False) def hook_before_train_loop(self): super().hook_before_train_loop() if self.is_caching_text_embeddings: # make sure model is on cpu for this part so we don't oom. self.sd.unet.to('cpu') # cache unconditional embeds (blank prompt) with torch.no_grad(): self.unconditional_embeds = self.sd.encode_prompt( [self.train_config.unconditional_prompt], long_prompts=self.do_long_prompts ).to( self.device_torch, dtype=self.sd.torch_dtype ).detach() if self.train_config.do_prior_divergence: self.do_prior_prediction = True # move vae to device if we did not cache latents if not self.is_latents_cached: self.sd.vae.eval() self.sd.vae.to(self.device_torch) else: # offload it. Already cached self.sd.vae.to('cpu') flush() add_all_snr_to_noise_scheduler(self.sd.noise_scheduler, self.device_torch) if self.adapter is not None: self.adapter.to(self.device_torch) # check if we have regs and using adapter and caching clip embeddings has_reg = self.datasets_reg is not None and len(self.datasets_reg) > 0 is_caching_clip_embeddings = self.datasets is not None and any([self.datasets[i].cache_clip_vision_to_disk for i in range(len(self.datasets))]) if has_reg and is_caching_clip_embeddings: # we need a list of unconditional clip image embeds from other datasets to handle regs unconditional_clip_image_embeds = [] datasets = get_dataloader_datasets(self.data_loader) for i in range(len(datasets)): unconditional_clip_image_embeds += datasets[i].clip_vision_unconditional_cache if len(unconditional_clip_image_embeds) == 0: raise ValueError("No unconditional clip image embeds found. This should not happen") self._clip_image_embeds_unconditional = unconditional_clip_image_embeds if self.train_config.negative_prompt is not None: if os.path.exists(self.train_config.negative_prompt): with open(self.train_config.negative_prompt, 'r') as f: self.negative_prompt_pool = f.readlines() # remove empty self.negative_prompt_pool = [x.strip() for x in self.negative_prompt_pool if x.strip() != ""] else: # single prompt self.negative_prompt_pool = [self.train_config.negative_prompt] # handle unload text encoder if self.train_config.unload_text_encoder or self.is_caching_text_embeddings: with torch.no_grad(): if self.train_config.train_text_encoder: raise ValueError("Cannot unload text encoder if training text encoder") # cache embeddings print_acc("\n***** UNLOADING TEXT ENCODER *****") if self.is_caching_text_embeddings: print_acc("Embeddings cached to disk. We dont need the text encoder anymore") else: print_acc("This will train only with a blank prompt or trigger word, if set") print_acc("If this is not what you want, remove the unload_text_encoder flag") print_acc("***********************************") print_acc("") self.sd.text_encoder_to(self.device_torch) self.cached_blank_embeds = self.sd.encode_prompt("") if self.trigger_word is not None: self.cached_trigger_embeds = self.sd.encode_prompt(self.trigger_word) if self.train_config.diff_output_preservation: self.diff_output_preservation_embeds = self.sd.encode_prompt(self.train_config.diff_output_preservation_class) self.cache_sample_prompts() # unload the text encoder if self.is_caching_text_embeddings: unload_text_encoder(self.sd) else: # todo once every model is tested to work, unload properly. Though, this will all be merged into one thing. # keep legacy usage for now. self.sd.text_encoder_to("cpu") flush() if self.train_config.diffusion_feature_extractor_path is not None: vae = self.sd.vae # if not (self.model_config.arch in ["flux"]) or self.sd.vae.__class__.__name__ == "AutoencoderPixelMixer": # vae = self.sd.vae self.dfe = load_dfe(self.train_config.diffusion_feature_extractor_path, vae=vae) self.dfe.to(self.device_torch) if hasattr(self.dfe, 'vision_encoder') and self.train_config.gradient_checkpointing: # must be set to train for gradient checkpointing to work self.dfe.vision_encoder.train() self.dfe.vision_encoder.gradient_checkpointing = True else: self.dfe.eval() # enable gradient checkpointing on the vae if vae is not None and self.train_config.gradient_checkpointing: vae.enable_gradient_checkpointing() vae.train() def process_output_for_turbo(self, pred, noisy_latents, timesteps, noise, batch): # to process turbo learning, we make one big step from our current timestep to the end # we then denoise the prediction on that remaining step and target our loss to our target latents # this currently only works on euler_a (that I know of). Would work on others, but needs to be coded to do so. # needs to be done on each item in batch as they may all have different timesteps batch_size = pred.shape[0] pred_chunks = torch.chunk(pred, batch_size, dim=0) noisy_latents_chunks = torch.chunk(noisy_latents, batch_size, dim=0) timesteps_chunks = torch.chunk(timesteps, batch_size, dim=0) latent_chunks = torch.chunk(batch.latents, batch_size, dim=0) noise_chunks = torch.chunk(noise, batch_size, dim=0) with torch.no_grad(): # set the timesteps to 1000 so we can capture them to calculate the sigmas self.sd.noise_scheduler.set_timesteps( self.sd.noise_scheduler.config.num_train_timesteps, device=self.device_torch ) train_timesteps = self.sd.noise_scheduler.timesteps.clone().detach() train_sigmas = self.sd.noise_scheduler.sigmas.clone().detach() # set the scheduler to one timestep, we build the step and sigmas for each item in batch for the partial step self.sd.noise_scheduler.set_timesteps( 1, device=self.device_torch ) denoised_pred_chunks = [] target_pred_chunks = [] for i in range(batch_size): pred_item = pred_chunks[i] noisy_latents_item = noisy_latents_chunks[i] timesteps_item = timesteps_chunks[i] latents_item = latent_chunks[i] noise_item = noise_chunks[i] with torch.no_grad(): timestep_idx = [(train_timesteps == t).nonzero().item() for t in timesteps_item][0] single_step_timestep_schedule = [timesteps_item.squeeze().item()] # extract the sigma idx for our midpoint timestep sigmas = train_sigmas[timestep_idx:timestep_idx + 1].to(self.device_torch) end_sigma_idx = random.randint(timestep_idx, len(train_sigmas) - 1) end_sigma = train_sigmas[end_sigma_idx:end_sigma_idx + 1].to(self.device_torch) # add noise to our target # build the big sigma step. The to step will now be to 0 giving it a full remaining denoising half step # self.sd.noise_scheduler.sigmas = torch.cat([sigmas, torch.zeros_like(sigmas)]).detach() self.sd.noise_scheduler.sigmas = torch.cat([sigmas, end_sigma]).detach() # set our single timstep self.sd.noise_scheduler.timesteps = torch.from_numpy( np.array(single_step_timestep_schedule, dtype=np.float32) ).to(device=self.device_torch) # set the step index to None so it will be recalculated on first step self.sd.noise_scheduler._step_index = None denoised_latent = self.sd.noise_scheduler.step( pred_item, timesteps_item, noisy_latents_item.detach(), return_dict=False )[0] residual_noise = (noise_item * end_sigma.flatten()).detach().to(self.device_torch, dtype=get_torch_dtype( self.train_config.dtype)) # remove the residual noise from the denoised latents. Output should be a clean prediction (theoretically) denoised_latent = denoised_latent - residual_noise denoised_pred_chunks.append(denoised_latent) denoised_latents = torch.cat(denoised_pred_chunks, dim=0) # set the scheduler back to the original timesteps self.sd.noise_scheduler.set_timesteps( self.sd.noise_scheduler.config.num_train_timesteps, device=self.device_torch ) output = denoised_latents / self.sd.vae.config['scaling_factor'] output = self.sd.vae.decode(output).sample if self.train_config.show_turbo_outputs: # since we are completely denoising, we can show them here with torch.no_grad(): show_tensors(output) # we return our big partial step denoised latents as our pred and our untouched latents as our target. # you can do mse against the two here or run the denoised through the vae for pixel space loss against the # input tensor images. return output, batch.tensor.to(self.device_torch, dtype=get_torch_dtype(self.train_config.dtype)) # you can expand these in a child class to make customization easier def calculate_loss( self, noise_pred: torch.Tensor, noise: torch.Tensor, noisy_latents: torch.Tensor, timesteps: torch.Tensor, batch: 'DataLoaderBatchDTO', mask_multiplier: Union[torch.Tensor, float] = 1.0, prior_pred: Union[torch.Tensor, None] = None, **kwargs ): loss_target = self.train_config.loss_target is_reg = any(batch.get_is_reg_list()) additional_loss = 0.0 prior_mask_multiplier = None target_mask_multiplier = None dtype = get_torch_dtype(self.train_config.dtype) has_mask = batch.mask_tensor is not None with torch.no_grad(): loss_multiplier = torch.tensor(batch.loss_multiplier_list).to(self.device_torch, dtype=torch.float32) if self.train_config.match_noise_norm: # match the norm of the noise noise_norm = torch.linalg.vector_norm(noise, ord=2, dim=(1, 2, 3), keepdim=True) noise_pred_norm = torch.linalg.vector_norm(noise_pred, ord=2, dim=(1, 2, 3), keepdim=True) noise_pred = noise_pred * (noise_norm / noise_pred_norm) if self.train_config.pred_scaler != 1.0: noise_pred = noise_pred * self.train_config.pred_scaler target = None if self.train_config.target_noise_multiplier != 1.0: noise = noise * self.train_config.target_noise_multiplier if self.train_config.correct_pred_norm or (self.train_config.inverted_mask_prior and prior_pred is not None and has_mask): if self.train_config.correct_pred_norm and not is_reg: with torch.no_grad(): # this only works if doing a prior pred if prior_pred is not None: prior_mean = prior_pred.mean([2,3], keepdim=True) prior_std = prior_pred.std([2,3], keepdim=True) noise_mean = noise_pred.mean([2,3], keepdim=True) noise_std = noise_pred.std([2,3], keepdim=True) mean_adjust = prior_mean - noise_mean std_adjust = prior_std - noise_std mean_adjust = mean_adjust * self.train_config.correct_pred_norm_multiplier std_adjust = std_adjust * self.train_config.correct_pred_norm_multiplier target_mean = noise_mean + mean_adjust target_std = noise_std + std_adjust eps = 1e-5 # match the noise to the prior noise = (noise - noise_mean) / (noise_std + eps) noise = noise * (target_std + eps) + target_mean noise = noise.detach() if self.train_config.inverted_mask_prior and prior_pred is not None and has_mask: assert not self.train_config.train_turbo with torch.no_grad(): prior_mask = batch.mask_tensor.to(self.device_torch, dtype=dtype) # resize to size of noise_pred prior_mask = torch.nn.functional.interpolate(prior_mask, size=(noise_pred.shape[2], noise_pred.shape[3]), mode='bicubic') # stack first channel to match channels of noise_pred prior_mask = torch.cat([prior_mask[:1]] * noise_pred.shape[1], dim=1) prior_mask_multiplier = 1.0 - prior_mask # scale so it is a mean of 1 prior_mask_multiplier = prior_mask_multiplier / prior_mask_multiplier.mean() if self.sd.is_flow_matching: target = (noise - batch.latents).detach() else: target = noise elif prior_pred is not None and not self.train_config.do_prior_divergence: assert not self.train_config.train_turbo # matching adapter prediction target = prior_pred elif self.sd.prediction_type == 'v_prediction': # v-parameterization training target = self.sd.noise_scheduler.get_velocity(batch.tensor, noise, timesteps) elif hasattr(self.sd, 'get_loss_target'): target = self.sd.get_loss_target( noise=noise, batch=batch, timesteps=timesteps, ).detach() elif self.sd.is_flow_matching: # forward ODE target = (noise - batch.latents).detach() # reverse ODE # target = (batch.latents - noise).detach() else: target = noise if self.dfe is not None: if self.dfe.version == 1: model = self.sd if model is not None and hasattr(model, 'get_stepped_pred'): stepped_latents = model.get_stepped_pred(noise_pred, noise) else: # stepped_latents = noise - noise_pred # first we step the scheduler from current timestep to the very end for a full denoise bs = noise_pred.shape[0] noise_pred_chunks = torch.chunk(noise_pred, bs) timestep_chunks = torch.chunk(timesteps, bs) noisy_latent_chunks = torch.chunk(noisy_latents, bs) stepped_chunks = [] for idx in range(bs): model_output = noise_pred_chunks[idx] timestep = timestep_chunks[idx] self.sd.noise_scheduler._step_index = None self.sd.noise_scheduler._init_step_index(timestep) sample = noisy_latent_chunks[idx].to(torch.float32) sigma = self.sd.noise_scheduler.sigmas[self.sd.noise_scheduler.step_index] sigma_next = self.sd.noise_scheduler.sigmas[-1] # use last sigma for final step prev_sample = sample + (sigma_next - sigma) * model_output stepped_chunks.append(prev_sample) stepped_latents = torch.cat(stepped_chunks, dim=0) stepped_latents = stepped_latents.to(self.sd.vae.device, dtype=self.sd.vae.dtype) # resize to half the size of the latents stepped_latents_half = torch.nn.functional.interpolate( stepped_latents, size=(stepped_latents.shape[2] // 2, stepped_latents.shape[3] // 2), mode='bilinear', align_corners=False ) pred_features = self.dfe(stepped_latents.float()) pred_features_half = self.dfe(stepped_latents_half.float()) with torch.no_grad(): target_features = self.dfe(batch.latents.to(self.device_torch, dtype=torch.float32)) batch_latents_half = torch.nn.functional.interpolate( batch.latents.to(self.device_torch, dtype=torch.float32), size=(batch.latents.shape[2] // 2, batch.latents.shape[3] // 2), mode='bilinear', align_corners=False ) target_features_half = self.dfe(batch_latents_half) # scale dfe so it is weaker at higher noise levels dfe_scaler = 1 - (timesteps.float() / 1000.0).view(-1, 1, 1, 1).to(self.device_torch) dfe_loss = torch.nn.functional.mse_loss(pred_features, target_features, reduction="none") * \ self.train_config.diffusion_feature_extractor_weight * dfe_scaler dfe_loss_half = torch.nn.functional.mse_loss(pred_features_half, target_features_half, reduction="none") * \ self.train_config.diffusion_feature_extractor_weight * dfe_scaler additional_loss += dfe_loss.mean() + dfe_loss_half.mean() elif self.dfe.version == 2: # version 2 # do diffusion feature extraction on target with torch.no_grad(): rectified_flow_target = noise.float() - batch.latents.float() target_feature_list = self.dfe(torch.cat([rectified_flow_target, noise.float()], dim=1)) # do diffusion feature extraction on prediction pred_feature_list = self.dfe(torch.cat([noise_pred.float(), noise.float()], dim=1)) dfe_loss = 0.0 for i in range(len(target_feature_list)): dfe_loss += torch.nn.functional.mse_loss(pred_feature_list[i], target_feature_list[i], reduction="mean") additional_loss += dfe_loss * self.train_config.diffusion_feature_extractor_weight * 100.0 elif self.dfe.version == 3 or self.dfe.version == 4: dfe_loss = self.dfe( noise=noise, noise_pred=noise_pred, noisy_latents=noisy_latents, timesteps=timesteps, batch=batch, scheduler=self.sd.noise_scheduler ) additional_loss += dfe_loss * self.train_config.diffusion_feature_extractor_weight else: raise ValueError(f"Unknown diffusion feature extractor version {self.dfe.version}") if self.train_config.do_guidance_loss: with torch.no_grad(): # we make cached blank prompt embeds that match the batch size unconditional_embeds = concat_prompt_embeds( [self.unconditional_embeds] * noisy_latents.shape[0], ) cfm_pred = self.predict_noise( noisy_latents=noisy_latents, timesteps=timesteps, conditional_embeds=unconditional_embeds, unconditional_embeds=None, batch=batch, ) # zero cfg # ref https://github.com/WeichenFan/CFG-Zero-star/blob/cdac25559e3f16cb95f0016c04c709ea1ab9452b/wan_pipeline.py#L557 batch_size = target.shape[0] positive_flat = target.view(batch_size, -1) negative_flat = cfm_pred.view(batch_size, -1) # Calculate dot production dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) # Squared norm of uncondition squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 # st_star = v_cond^T * v_uncond / ||v_uncond||^2 st_star = dot_product / squared_norm alpha = st_star is_video = len(target.shape) == 5 alpha = alpha.view(batch_size, 1, 1, 1) if not is_video else alpha.view(batch_size, 1, 1, 1, 1) guidance_scale = self._guidance_loss_target_batch 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) unconditional_target = cfm_pred * alpha target = unconditional_target + guidance_scale * (target - unconditional_target) if target is None: target = noise pred = noise_pred if self.train_config.train_turbo: pred, target = self.process_output_for_turbo(pred, noisy_latents, timesteps, noise, batch) ignore_snr = False if loss_target == 'source' or loss_target == 'unaugmented': assert not self.train_config.train_turbo # ignore_snr = True if batch.sigmas is None: raise ValueError("Batch sigmas is None. This should not happen") # src https://github.com/huggingface/diffusers/blob/324d18fba23f6c9d7475b0ff7c777685f7128d40/examples/t2i_adapter/train_t2i_adapter_sdxl.py#L1190 denoised_latents = noise_pred * (-batch.sigmas) + noisy_latents weighing = batch.sigmas ** -2.0 if loss_target == 'source': # denoise the latent and compare to the latent in the batch target = batch.latents elif loss_target == 'unaugmented': # we have to encode images into latents for now # we also denoise as the unaugmented tensor is not a noisy diffirental with torch.no_grad(): unaugmented_latents = self.sd.encode_images(batch.unaugmented_tensor).to(self.device_torch, dtype=dtype) unaugmented_latents = unaugmented_latents * self.train_config.latent_multiplier target = unaugmented_latents.detach() # Get the target for loss depending on the prediction type if self.sd.noise_scheduler.config.prediction_type == "epsilon": target = target # we are computing loss against denoise latents elif self.sd.noise_scheduler.config.prediction_type == "v_prediction": target = self.sd.noise_scheduler.get_velocity(target, noise, timesteps) else: raise ValueError(f"Unknown prediction type {self.sd.noise_scheduler.config.prediction_type}") # mse loss without reduction loss_per_element = (weighing.float() * (denoised_latents.float() - target.float()) ** 2) loss = loss_per_element else: if self.train_config.loss_type == "mae": loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none") elif self.train_config.loss_type == "wavelet": loss = wavelet_loss(pred, batch.latents, noise) else: loss = torch.nn.functional.mse_loss(pred.float(), target.float(), reduction="none") do_weighted_timesteps = False if self.sd.is_flow_matching: if self.train_config.linear_timesteps or self.train_config.linear_timesteps2: do_weighted_timesteps = True if self.train_config.timestep_type == "weighted": # use the noise scheduler to get the weights for the timesteps do_weighted_timesteps = True # handle linear timesteps and only adjust the weight of the timesteps if do_weighted_timesteps: # calculate the weights for the timesteps timestep_weight = self.sd.noise_scheduler.get_weights_for_timesteps( timesteps, v2=self.train_config.linear_timesteps2, timestep_type=self.train_config.timestep_type ).to(loss.device, dtype=loss.dtype) if len(loss.shape) == 4: timestep_weight = timestep_weight.view(-1, 1, 1, 1).detach() elif len(loss.shape) == 5: timestep_weight = timestep_weight.view(-1, 1, 1, 1, 1).detach() loss = loss * timestep_weight if self.train_config.do_prior_divergence and prior_pred is not None: loss = loss + (torch.nn.functional.mse_loss(pred.float(), prior_pred.float(), reduction="none") * -1.0) if self.train_config.train_turbo: mask_multiplier = mask_multiplier[:, 3:, :, :] # resize to the size of the loss mask_multiplier = torch.nn.functional.interpolate(mask_multiplier, size=(pred.shape[2], pred.shape[3]), mode='nearest') # multiply by our mask try: loss = loss * mask_multiplier except: # todo handle mask with video models pass prior_loss = None if self.train_config.inverted_mask_prior and prior_pred is not None and prior_mask_multiplier is not None: assert not self.train_config.train_turbo if self.train_config.loss_type == "mae": prior_loss = torch.nn.functional.l1_loss(pred.float(), prior_pred.float(), reduction="none") else: prior_loss = torch.nn.functional.mse_loss(pred.float(), prior_pred.float(), reduction="none") prior_loss = prior_loss * prior_mask_multiplier * self.train_config.inverted_mask_prior_multiplier if torch.isnan(prior_loss).any(): print_acc("Prior loss is nan") prior_loss = None else: prior_loss = prior_loss.mean([1, 2, 3]) # loss = loss + prior_loss # loss = loss + prior_loss # loss = loss + prior_loss loss = loss.mean([1, 2, 3]) # apply loss multiplier before prior loss # multiply by our mask try: loss = loss * loss_multiplier except: # todo handle mask with video models pass if prior_loss is not None: loss = loss + prior_loss if not self.train_config.train_turbo: if self.train_config.learnable_snr_gos: # add snr_gamma loss = apply_learnable_snr_gos(loss, timesteps, self.snr_gos) elif self.train_config.snr_gamma is not None and self.train_config.snr_gamma > 0.000001 and not ignore_snr: # add snr_gamma loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.snr_gamma, fixed=True) elif self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001 and not ignore_snr: # add min_snr_gamma loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.min_snr_gamma) loss = loss.mean() # check for additional losses if self.adapter is not None and hasattr(self.adapter, "additional_loss") and self.adapter.additional_loss is not None: loss = loss + self.adapter.additional_loss.mean() self.adapter.additional_loss = None if self.train_config.target_norm_std: # seperate out the batch and channels pred_std = noise_pred.std([2, 3], keepdim=True) norm_std_loss = torch.abs(self.train_config.target_norm_std_value - pred_std).mean() loss = loss + norm_std_loss return loss + additional_loss def preprocess_batch(self, batch: 'DataLoaderBatchDTO'): return batch def get_guided_loss( self, noisy_latents: torch.Tensor, conditional_embeds: PromptEmbeds, match_adapter_assist: bool, network_weight_list: list, timesteps: torch.Tensor, pred_kwargs: dict, batch: 'DataLoaderBatchDTO', noise: torch.Tensor, unconditional_embeds: Optional[PromptEmbeds] = None, **kwargs ): loss = get_guidance_loss( noisy_latents=noisy_latents, conditional_embeds=conditional_embeds, match_adapter_assist=match_adapter_assist, network_weight_list=network_weight_list, timesteps=timesteps, pred_kwargs=pred_kwargs, batch=batch, noise=noise, sd=self.sd, unconditional_embeds=unconditional_embeds, train_config=self.train_config, **kwargs ) return loss # ------------------------------------------------------------------ # Mean-Flow loss (Geng et al., “Mean Flows for One-step Generative # Modelling”, 2025 – see Alg. 1 + Eq. (6) of the paper) # This version avoids jvp / double-back-prop issues with Flash-Attention # adapted from the work of lodestonerock # ------------------------------------------------------------------ def get_mean_flow_loss( self, noisy_latents: torch.Tensor, conditional_embeds: PromptEmbeds, match_adapter_assist: bool, network_weight_list: list, timesteps: torch.Tensor, pred_kwargs: dict, batch: 'DataLoaderBatchDTO', noise: torch.Tensor, unconditional_embeds: Optional[PromptEmbeds] = None, **kwargs ): dtype = get_torch_dtype(self.train_config.dtype) total_steps = float(self.sd.noise_scheduler.config.num_train_timesteps) # e.g. 1000 base_eps = 1e-3 min_time_gap = 1e-2 with torch.no_grad(): num_train_timesteps = self.sd.noise_scheduler.config.num_train_timesteps batch_size = batch.latents.shape[0] timestep_t_list = [] timestep_r_list = [] for i in range(batch_size): t1 = random.randint(0, num_train_timesteps - 1) t2 = random.randint(0, num_train_timesteps - 1) t_t = self.sd.noise_scheduler.timesteps[min(t1, t2)] t_r = self.sd.noise_scheduler.timesteps[max(t1, t2)] if (t_t - t_r).item() < min_time_gap * 1000: scaled_time_gap = min_time_gap * 1000 if t_t.item() + scaled_time_gap > 1000: t_r = t_r - scaled_time_gap else: t_t = t_t + scaled_time_gap timestep_t_list.append(t_t) timestep_r_list.append(t_r) timesteps_t = torch.stack(timestep_t_list, dim=0).float() timesteps_r = torch.stack(timestep_r_list, dim=0).float() t_frac = timesteps_t / total_steps # [0,1] r_frac = timesteps_r / total_steps # [0,1] latents_clean = batch.latents.to(dtype) noise_sample = noise.to(dtype) lerp_vector = latents_clean * (1.0 - t_frac[:, None, None, None]) + noise_sample * t_frac[:, None, None, None] eps = base_eps # concatenate timesteps as input for u(z, r, t) timesteps_cat = torch.cat([t_frac, r_frac], dim=0) * total_steps # model predicts u(z, r, t) u_pred = self.predict_noise( noisy_latents=lerp_vector.to(dtype), timesteps=timesteps_cat.to(dtype), conditional_embeds=conditional_embeds, unconditional_embeds=unconditional_embeds, batch=batch, **pred_kwargs ) with torch.no_grad(): t_frac_plus_eps = (t_frac + eps).clamp(0.0, 1.0) lerp_perturbed = latents_clean * (1.0 - t_frac_plus_eps[:, None, None, None]) + noise_sample * t_frac_plus_eps[:, None, None, None] timesteps_cat_perturbed = torch.cat([t_frac_plus_eps, r_frac], dim=0) * total_steps u_perturbed = self.predict_noise( noisy_latents=lerp_perturbed.to(dtype), timesteps=timesteps_cat_perturbed.to(dtype), conditional_embeds=conditional_embeds, unconditional_embeds=unconditional_embeds, batch=batch, **pred_kwargs ) # compute du/dt via finite difference (detached) du_dt = (u_perturbed - u_pred).detach() / eps # du_dt = (u_perturbed - u_pred).detach() du_dt = du_dt.to(dtype) time_gap = (t_frac - r_frac)[:, None, None, None].to(dtype) time_gap.clamp(min=1e-4) u_shifted = u_pred + time_gap * du_dt # u_shifted = u_pred + du_dt / time_gap # u_shifted = u_pred # a step is done like this: # stepped_latent = model_input + (timestep_next - timestep) * model_output # flow target velocity # v_target = (noise_sample - latents_clean) / time_gap # flux predicts opposite of velocity, so we need to invert it v_target = (latents_clean - noise_sample) / time_gap # compute loss loss = torch.nn.functional.mse_loss( u_shifted.float(), v_target.float(), reduction='none' ) with torch.no_grad(): pure_loss = loss.mean().detach() pure_loss.requires_grad_(True) loss = loss.mean() if loss.item() > 1e3: pass self.accelerator.backward(loss) return pure_loss def get_prior_prediction( self, noisy_latents: torch.Tensor, conditional_embeds: PromptEmbeds, match_adapter_assist: bool, network_weight_list: list, timesteps: torch.Tensor, pred_kwargs: dict, batch: 'DataLoaderBatchDTO', noise: torch.Tensor, unconditional_embeds: Optional[PromptEmbeds] = None, conditioned_prompts=None, **kwargs ): # todo for embeddings, we need to run without trigger words was_unet_training = self.sd.unet.training was_network_active = False if self.network is not None: was_network_active = self.network.is_active self.network.is_active = False can_disable_adapter = False was_adapter_active = False if self.adapter is not None and (isinstance(self.adapter, IPAdapter) or isinstance(self.adapter, ReferenceAdapter) or (isinstance(self.adapter, CustomAdapter)) ): can_disable_adapter = True was_adapter_active = self.adapter.is_active self.adapter.is_active = False if self.train_config.unload_text_encoder and self.adapter is not None and not isinstance(self.adapter, CustomAdapter): raise ValueError("Prior predictions currently do not support unloading text encoder with adapter") # do a prediction here so we can match its output with network multiplier set to 0.0 with torch.no_grad(): dtype = get_torch_dtype(self.train_config.dtype) embeds_to_use = conditional_embeds.clone().detach() # handle clip vision adapter by removing triggers from prompt and replacing with the class name if (self.adapter is not None and isinstance(self.adapter, ClipVisionAdapter)) or self.embedding is not None: prompt_list = batch.get_caption_list() class_name = '' triggers = ['[trigger]', '[name]'] remove_tokens = [] if self.embed_config is not None: triggers.append(self.embed_config.trigger) for i in range(1, self.embed_config.tokens): remove_tokens.append(f"{self.embed_config.trigger}_{i}") if self.embed_config.trigger_class_name is not None: class_name = self.embed_config.trigger_class_name if self.adapter is not None: triggers.append(self.adapter_config.trigger) for i in range(1, self.adapter_config.num_tokens): remove_tokens.append(f"{self.adapter_config.trigger}_{i}") if self.adapter_config.trigger_class_name is not None: class_name = self.adapter_config.trigger_class_name for idx, prompt in enumerate(prompt_list): for remove_token in remove_tokens: prompt = prompt.replace(remove_token, '') for trigger in triggers: prompt = prompt.replace(trigger, class_name) prompt_list[idx] = prompt if batch.prompt_embeds is not None: embeds_to_use = batch.prompt_embeds.clone().to(self.device_torch, dtype=dtype) else: embeds_to_use = self.sd.encode_prompt( prompt_list, long_prompts=self.do_long_prompts).to( self.device_torch, dtype=dtype).detach() # dont use network on this # self.network.multiplier = 0.0 self.sd.unet.eval() if self.adapter is not None and isinstance(self.adapter, IPAdapter) and not self.sd.is_flux and not self.sd.is_lumina2: # we need to remove the image embeds from the prompt except for flux embeds_to_use: PromptEmbeds = embeds_to_use.clone().detach() end_pos = embeds_to_use.text_embeds.shape[1] - self.adapter_config.num_tokens embeds_to_use.text_embeds = embeds_to_use.text_embeds[:, :end_pos, :] if unconditional_embeds is not None: unconditional_embeds = unconditional_embeds.clone().detach() unconditional_embeds.text_embeds = unconditional_embeds.text_embeds[:, :end_pos] if unconditional_embeds is not None: unconditional_embeds = unconditional_embeds.to(self.device_torch, dtype=dtype).detach() guidance_embedding_scale = self.train_config.cfg_scale if self.train_config.do_guidance_loss: guidance_embedding_scale = self._guidance_loss_target_batch prior_pred = self.sd.predict_noise( latents=noisy_latents.to(self.device_torch, dtype=dtype).detach(), conditional_embeddings=embeds_to_use.to(self.device_torch, dtype=dtype).detach(), unconditional_embeddings=unconditional_embeds, timestep=timesteps, guidance_scale=self.train_config.cfg_scale, guidance_embedding_scale=guidance_embedding_scale, rescale_cfg=self.train_config.cfg_rescale, batch=batch, **pred_kwargs # adapter residuals in here ) if was_unet_training: self.sd.unet.train() prior_pred = prior_pred.detach() # remove the residuals as we wont use them on prediction when matching control if match_adapter_assist and 'down_intrablock_additional_residuals' in pred_kwargs: del pred_kwargs['down_intrablock_additional_residuals'] if match_adapter_assist and 'down_block_additional_residuals' in pred_kwargs: del pred_kwargs['down_block_additional_residuals'] if match_adapter_assist and 'mid_block_additional_residual' in pred_kwargs: del pred_kwargs['mid_block_additional_residual'] if can_disable_adapter: self.adapter.is_active = was_adapter_active # restore network # self.network.multiplier = network_weight_list if self.network is not None: self.network.is_active = was_network_active return prior_pred def before_unet_predict(self): pass def after_unet_predict(self): pass def end_of_training_loop(self): pass def predict_noise( self, noisy_latents: torch.Tensor, timesteps: Union[int, torch.Tensor] = 1, conditional_embeds: Union[PromptEmbeds, None] = None, unconditional_embeds: Union[PromptEmbeds, None] = None, batch: Optional['DataLoaderBatchDTO'] = None, is_primary_pred: bool = False, **kwargs, ): dtype = get_torch_dtype(self.train_config.dtype) guidance_embedding_scale = self.train_config.cfg_scale if self.train_config.do_guidance_loss: guidance_embedding_scale = self._guidance_loss_target_batch return self.sd.predict_noise( latents=noisy_latents.to(self.device_torch, dtype=dtype), conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype), unconditional_embeddings=unconditional_embeds, timestep=timesteps, guidance_scale=self.train_config.cfg_scale, guidance_embedding_scale=guidance_embedding_scale, detach_unconditional=False, rescale_cfg=self.train_config.cfg_rescale, bypass_guidance_embedding=self.train_config.bypass_guidance_embedding, batch=batch, **kwargs ) def train_single_accumulation(self, batch: DataLoaderBatchDTO): with torch.no_grad(): self.timer.start('preprocess_batch') if isinstance(self.adapter, CustomAdapter): batch = self.adapter.edit_batch_raw(batch) batch = self.preprocess_batch(batch) if isinstance(self.adapter, CustomAdapter): batch = self.adapter.edit_batch_processed(batch) dtype = get_torch_dtype(self.train_config.dtype) # sanity check if self.sd.vae.dtype != self.sd.vae_torch_dtype: self.sd.vae = self.sd.vae.to(self.sd.vae_torch_dtype) if isinstance(self.sd.text_encoder, list): for encoder in self.sd.text_encoder: if encoder.dtype != self.sd.te_torch_dtype: encoder.to(self.sd.te_torch_dtype) else: if self.sd.text_encoder.dtype != self.sd.te_torch_dtype: self.sd.text_encoder.to(self.sd.te_torch_dtype) noisy_latents, noise, timesteps, conditioned_prompts, imgs = self.process_general_training_batch(batch) if self.train_config.do_cfg or self.train_config.do_random_cfg: # pick random negative prompts if self.negative_prompt_pool is not None: negative_prompts = [] for i in range(noisy_latents.shape[0]): num_neg = random.randint(1, self.train_config.max_negative_prompts) this_neg_prompts = [random.choice(self.negative_prompt_pool) for _ in range(num_neg)] this_neg_prompt = ', '.join(this_neg_prompts) negative_prompts.append(this_neg_prompt) self.batch_negative_prompt = negative_prompts else: self.batch_negative_prompt = ['' for _ in range(batch.latents.shape[0])] if self.adapter and isinstance(self.adapter, CustomAdapter): # condition the prompt # todo handle more than one adapter image conditioned_prompts = self.adapter.condition_prompt(conditioned_prompts) network_weight_list = batch.get_network_weight_list() if self.train_config.single_item_batching: network_weight_list = network_weight_list + network_weight_list has_adapter_img = batch.control_tensor is not None has_clip_image = batch.clip_image_tensor is not None has_clip_image_embeds = batch.clip_image_embeds is not None # force it to be true if doing regs as we handle those differently if any([batch.file_items[idx].is_reg for idx in range(len(batch.file_items))]): has_clip_image = True if self._clip_image_embeds_unconditional is not None: has_clip_image_embeds = True # we are caching embeds, handle that differently has_clip_image = False if self.adapter is not None and isinstance(self.adapter, IPAdapter) and not has_clip_image and has_adapter_img: raise ValueError( "IPAdapter control image is now 'clip_image_path' instead of 'control_path'. Please update your dataset config ") match_adapter_assist = False # check if we are matching the adapter assistant if self.assistant_adapter: if self.train_config.match_adapter_chance == 1.0: match_adapter_assist = True elif self.train_config.match_adapter_chance > 0.0: match_adapter_assist = torch.rand( (1,), device=self.device_torch, dtype=dtype ) < self.train_config.match_adapter_chance self.timer.stop('preprocess_batch') is_reg = False loss_multiplier = torch.ones((noisy_latents.shape[0], 1, 1, 1), device=self.device_torch, dtype=dtype) for idx, file_item in enumerate(batch.file_items): if file_item.is_reg: loss_multiplier[idx] = loss_multiplier[idx] * self.train_config.reg_weight is_reg = True adapter_images = None sigmas = None if has_adapter_img and (self.adapter or self.assistant_adapter): with self.timer('get_adapter_images'): # todo move this to data loader if batch.control_tensor is not None: adapter_images = batch.control_tensor.to(self.device_torch, dtype=dtype).detach() # match in channels if self.assistant_adapter is not None: in_channels = self.assistant_adapter.config.in_channels if adapter_images.shape[1] != in_channels: # we need to match the channels adapter_images = adapter_images[:, :in_channels, :, :] else: raise NotImplementedError("Adapter images now must be loaded with dataloader") clip_images = None if has_clip_image: with self.timer('get_clip_images'): # todo move this to data loader if batch.clip_image_tensor is not None: clip_images = batch.clip_image_tensor.to(self.device_torch, dtype=dtype).detach() mask_multiplier = torch.ones((noisy_latents.shape[0], 1, 1, 1), device=self.device_torch, dtype=dtype) if batch.mask_tensor is not None: with self.timer('get_mask_multiplier'): # upsampling no supported for bfloat16 mask_multiplier = batch.mask_tensor.to(self.device_torch, dtype=torch.float16).detach() # scale down to the size of the latents, mask multiplier shape(bs, 1, width, height), noisy_latents shape(bs, channels, width, height) mask_multiplier = torch.nn.functional.interpolate( mask_multiplier, size=(noisy_latents.shape[2], noisy_latents.shape[3]) ) # expand to match latents mask_multiplier = mask_multiplier.expand(-1, noisy_latents.shape[1], -1, -1) mask_multiplier = mask_multiplier.to(self.device_torch, dtype=dtype).detach() # make avg 1.0 mask_multiplier = mask_multiplier / mask_multiplier.mean() def get_adapter_multiplier(): if self.adapter and isinstance(self.adapter, T2IAdapter): # training a t2i adapter, not using as assistant. return 1.0 elif match_adapter_assist: # training a texture. We want it high adapter_strength_min = 0.9 adapter_strength_max = 1.0 else: # training with assistance, we want it low # adapter_strength_min = 0.4 # adapter_strength_max = 0.7 adapter_strength_min = 0.5 adapter_strength_max = 1.1 adapter_conditioning_scale = torch.rand( (1,), device=self.device_torch, dtype=dtype ) adapter_conditioning_scale = value_map( adapter_conditioning_scale, 0.0, 1.0, adapter_strength_min, adapter_strength_max ) return adapter_conditioning_scale # flush() with self.timer('grad_setup'): # text encoding grad_on_text_encoder = False if self.train_config.train_text_encoder: grad_on_text_encoder = True if self.embedding is not None: grad_on_text_encoder = True if self.adapter and isinstance(self.adapter, ClipVisionAdapter): grad_on_text_encoder = True if self.adapter_config and self.adapter_config.type == 'te_augmenter': grad_on_text_encoder = True # have a blank network so we can wrap it in a context and set multipliers without checking every time if self.network is not None: network = self.network else: network = BlankNetwork() # set the weights network.multiplier = network_weight_list # activate network if it exits prompts_1 = conditioned_prompts prompts_2 = None if self.train_config.short_and_long_captions_encoder_split and self.sd.is_xl: prompts_1 = batch.get_caption_short_list() prompts_2 = conditioned_prompts # make the batch splits if self.train_config.single_item_batching: if self.model_config.refiner_name_or_path is not None: raise ValueError("Single item batching is not supported when training the refiner") batch_size = noisy_latents.shape[0] # chunk/split everything noisy_latents_list = torch.chunk(noisy_latents, batch_size, dim=0) noise_list = torch.chunk(noise, batch_size, dim=0) timesteps_list = torch.chunk(timesteps, batch_size, dim=0) conditioned_prompts_list = [[prompt] for prompt in prompts_1] if imgs is not None: imgs_list = torch.chunk(imgs, batch_size, dim=0) else: imgs_list = [None for _ in range(batch_size)] if adapter_images is not None: adapter_images_list = torch.chunk(adapter_images, batch_size, dim=0) else: adapter_images_list = [None for _ in range(batch_size)] if clip_images is not None: clip_images_list = torch.chunk(clip_images, batch_size, dim=0) else: clip_images_list = [None for _ in range(batch_size)] mask_multiplier_list = torch.chunk(mask_multiplier, batch_size, dim=0) if prompts_2 is None: prompt_2_list = [None for _ in range(batch_size)] else: prompt_2_list = [[prompt] for prompt in prompts_2] else: noisy_latents_list = [noisy_latents] noise_list = [noise] timesteps_list = [timesteps] conditioned_prompts_list = [prompts_1] imgs_list = [imgs] adapter_images_list = [adapter_images] clip_images_list = [clip_images] mask_multiplier_list = [mask_multiplier] if prompts_2 is None: prompt_2_list = [None] else: prompt_2_list = [prompts_2] for noisy_latents, noise, timesteps, conditioned_prompts, imgs, adapter_images, clip_images, mask_multiplier, prompt_2 in zip( noisy_latents_list, noise_list, timesteps_list, conditioned_prompts_list, imgs_list, adapter_images_list, clip_images_list, mask_multiplier_list, prompt_2_list ): # if self.train_config.negative_prompt is not None: # # add negative prompt # conditioned_prompts = conditioned_prompts + [self.train_config.negative_prompt for x in # range(len(conditioned_prompts))] # if prompt_2 is not None: # prompt_2 = prompt_2 + [self.train_config.negative_prompt for x in range(len(prompt_2))] with (network): # encode clip adapter here so embeds are active for tokenizer if self.adapter and isinstance(self.adapter, ClipVisionAdapter): with self.timer('encode_clip_vision_embeds'): if has_clip_image: conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors( clip_images.detach().to(self.device_torch, dtype=dtype), is_training=True, has_been_preprocessed=True ) else: # just do a blank one conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors( torch.zeros( (noisy_latents.shape[0], 3, 512, 512), device=self.device_torch, dtype=dtype ), is_training=True, has_been_preprocessed=True, drop=True ) # it will be injected into the tokenizer when called self.adapter(conditional_clip_embeds) # do the custom adapter after the prior prediction if self.adapter and isinstance(self.adapter, CustomAdapter) and (has_clip_image or is_reg): quad_count = random.randint(1, 4) self.adapter.train() self.adapter.trigger_pre_te( tensors_preprocessed=clip_images if not is_reg else None, # on regs we send none to get random noise is_training=True, has_been_preprocessed=True, quad_count=quad_count, batch_tensor=batch.tensor if not is_reg else None, batch_size=noisy_latents.shape[0] ) with self.timer('encode_prompt'): unconditional_embeds = None if self.train_config.unload_text_encoder or self.is_caching_text_embeddings: with torch.set_grad_enabled(False): if batch.prompt_embeds is not None: # use the cached embeds conditional_embeds = batch.prompt_embeds.clone().detach().to( self.device_torch, dtype=dtype ) else: embeds_to_use = self.cached_blank_embeds.clone().detach().to( self.device_torch, dtype=dtype ) if self.cached_trigger_embeds is not None and not is_reg: embeds_to_use = self.cached_trigger_embeds.clone().detach().to( self.device_torch, dtype=dtype ) conditional_embeds = concat_prompt_embeds( [embeds_to_use] * noisy_latents.shape[0] ) if self.train_config.do_cfg: unconditional_embeds = self.cached_blank_embeds.clone().detach().to( self.device_torch, dtype=dtype ) unconditional_embeds = concat_prompt_embeds( [unconditional_embeds] * noisy_latents.shape[0] ) if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = False elif grad_on_text_encoder: with torch.set_grad_enabled(True): if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = False conditional_embeds = self.sd.encode_prompt( conditioned_prompts, prompt_2, dropout_prob=self.train_config.prompt_dropout_prob, long_prompts=self.do_long_prompts).to( self.device_torch, dtype=dtype) if self.train_config.do_cfg: if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = True # todo only do one and repeat it unconditional_embeds = self.sd.encode_prompt( self.batch_negative_prompt, self.batch_negative_prompt, dropout_prob=self.train_config.prompt_dropout_prob, long_prompts=self.do_long_prompts).to( self.device_torch, dtype=dtype) if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = False else: with torch.set_grad_enabled(False): # make sure it is in eval mode if isinstance(self.sd.text_encoder, list): for te in self.sd.text_encoder: te.eval() else: self.sd.text_encoder.eval() if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = False conditional_embeds = self.sd.encode_prompt( conditioned_prompts, prompt_2, dropout_prob=self.train_config.prompt_dropout_prob, long_prompts=self.do_long_prompts).to( self.device_torch, dtype=dtype) if self.train_config.do_cfg: if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = True unconditional_embeds = self.sd.encode_prompt( self.batch_negative_prompt, dropout_prob=self.train_config.prompt_dropout_prob, long_prompts=self.do_long_prompts).to( self.device_torch, dtype=dtype) if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = False if self.train_config.diff_output_preservation: dop_prompts = [p.replace(self.trigger_word, self.train_config.diff_output_preservation_class) for p in conditioned_prompts] dop_prompts_2 = None if prompt_2 is not None: dop_prompts_2 = [p.replace(self.trigger_word, self.train_config.diff_output_preservation_class) for p in prompt_2] self.diff_output_preservation_embeds = self.sd.encode_prompt( dop_prompts, dop_prompts_2, dropout_prob=self.train_config.prompt_dropout_prob, long_prompts=self.do_long_prompts).to( self.device_torch, dtype=dtype) # detach the embeddings conditional_embeds = conditional_embeds.detach() if self.train_config.do_cfg: unconditional_embeds = unconditional_embeds.detach() if self.decorator: conditional_embeds.text_embeds = self.decorator( conditional_embeds.text_embeds ) if self.train_config.do_cfg: unconditional_embeds.text_embeds = self.decorator( unconditional_embeds.text_embeds, is_unconditional=True ) # flush() pred_kwargs = {} if has_adapter_img: if (self.adapter and isinstance(self.adapter, T2IAdapter)) or ( self.assistant_adapter and isinstance(self.assistant_adapter, T2IAdapter)): with torch.set_grad_enabled(self.adapter is not None): adapter = self.assistant_adapter if self.assistant_adapter is not None else self.adapter adapter_multiplier = get_adapter_multiplier() with self.timer('encode_adapter'): down_block_additional_residuals = adapter(adapter_images) if self.assistant_adapter: # not training. detach down_block_additional_residuals = [ sample.to(dtype=dtype).detach() * adapter_multiplier for sample in down_block_additional_residuals ] else: down_block_additional_residuals = [ sample.to(dtype=dtype) * adapter_multiplier for sample in down_block_additional_residuals ] pred_kwargs['down_intrablock_additional_residuals'] = down_block_additional_residuals if self.adapter and isinstance(self.adapter, IPAdapter): with self.timer('encode_adapter_embeds'): # number of images to do if doing a quad image quad_count = random.randint(1, 4) image_size = self.adapter.input_size if has_clip_image_embeds: # todo handle reg images better than this if is_reg: # get unconditional image embeds from cache embeds = [ load_file(random.choice(batch.clip_image_embeds_unconditional)) for i in range(noisy_latents.shape[0]) ] conditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache( embeds, quad_count=quad_count ) if self.train_config.do_cfg: embeds = [ load_file(random.choice(batch.clip_image_embeds_unconditional)) for i in range(noisy_latents.shape[0]) ] unconditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache( embeds, quad_count=quad_count ) else: conditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache( batch.clip_image_embeds, quad_count=quad_count ) if self.train_config.do_cfg: unconditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache( batch.clip_image_embeds_unconditional, quad_count=quad_count ) elif is_reg: # we will zero it out in the img embedder clip_images = torch.zeros( (noisy_latents.shape[0], 3, image_size, image_size), device=self.device_torch, dtype=dtype ).detach() # drop will zero it out conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors( clip_images, drop=True, is_training=True, has_been_preprocessed=False, quad_count=quad_count ) if self.train_config.do_cfg: unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors( torch.zeros( (noisy_latents.shape[0], 3, image_size, image_size), device=self.device_torch, dtype=dtype ).detach(), is_training=True, drop=True, has_been_preprocessed=False, quad_count=quad_count ) elif has_clip_image: conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors( clip_images.detach().to(self.device_torch, dtype=dtype), is_training=True, has_been_preprocessed=True, quad_count=quad_count, # do cfg on clip embeds to normalize the embeddings for when doing cfg # cfg_embed_strength=3.0 if not self.train_config.do_cfg else None # cfg_embed_strength=3.0 if not self.train_config.do_cfg else None ) if self.train_config.do_cfg: unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors( clip_images.detach().to(self.device_torch, dtype=dtype), is_training=True, drop=True, has_been_preprocessed=True, quad_count=quad_count ) else: print_acc("No Clip Image") print_acc([file_item.path for file_item in batch.file_items]) raise ValueError("Could not find clip image") if not self.adapter_config.train_image_encoder: # we are not training the image encoder, so we need to detach the embeds conditional_clip_embeds = conditional_clip_embeds.detach() if self.train_config.do_cfg: unconditional_clip_embeds = unconditional_clip_embeds.detach() with self.timer('encode_adapter'): self.adapter.train() conditional_embeds = self.adapter( conditional_embeds.detach(), conditional_clip_embeds, is_unconditional=False ) if self.train_config.do_cfg: unconditional_embeds = self.adapter( unconditional_embeds.detach(), unconditional_clip_embeds, is_unconditional=True ) else: # wipe out unconsitional self.adapter.last_unconditional = None if self.adapter and isinstance(self.adapter, ReferenceAdapter): # pass in our scheduler self.adapter.noise_scheduler = self.lr_scheduler if has_clip_image or has_adapter_img: img_to_use = clip_images if has_clip_image else adapter_images # currently 0-1 needs to be -1 to 1 reference_images = ((img_to_use - 0.5) * 2).detach().to(self.device_torch, dtype=dtype) self.adapter.set_reference_images(reference_images) self.adapter.noise_scheduler = self.sd.noise_scheduler elif is_reg: self.adapter.set_blank_reference_images(noisy_latents.shape[0]) else: self.adapter.set_reference_images(None) prior_pred = None do_reg_prior = False # if is_reg and (self.network is not None or self.adapter is not None): # # we are doing a reg image and we have a network or adapter # do_reg_prior = True do_inverted_masked_prior = False if self.train_config.inverted_mask_prior and batch.mask_tensor is not None: do_inverted_masked_prior = True do_correct_pred_norm_prior = self.train_config.correct_pred_norm do_guidance_prior = False if batch.unconditional_latents is not None: # for this not that, we need a prior pred to normalize guidance_type: GuidanceType = batch.file_items[0].dataset_config.guidance_type if guidance_type == 'tnt': do_guidance_prior = True if (( has_adapter_img and self.assistant_adapter and match_adapter_assist) or self.do_prior_prediction or do_guidance_prior or do_reg_prior or do_inverted_masked_prior or self.train_config.correct_pred_norm): with self.timer('prior predict'): prior_embeds_to_use = conditional_embeds # use diff_output_preservation embeds if doing dfe if self.train_config.diff_output_preservation: prior_embeds_to_use = self.diff_output_preservation_embeds.expand_to_batch(noisy_latents.shape[0]) prior_pred = self.get_prior_prediction( noisy_latents=noisy_latents, conditional_embeds=prior_embeds_to_use, match_adapter_assist=match_adapter_assist, network_weight_list=network_weight_list, timesteps=timesteps, pred_kwargs=pred_kwargs, noise=noise, batch=batch, unconditional_embeds=unconditional_embeds, conditioned_prompts=conditioned_prompts ) if prior_pred is not None: prior_pred = prior_pred.detach() # do the custom adapter after the prior prediction if self.adapter and isinstance(self.adapter, CustomAdapter) and (has_clip_image or self.adapter_config.type in ['llm_adapter', 'text_encoder']): quad_count = random.randint(1, 4) self.adapter.train() conditional_embeds = self.adapter.condition_encoded_embeds( tensors_0_1=clip_images, prompt_embeds=conditional_embeds, is_training=True, has_been_preprocessed=True, quad_count=quad_count ) if self.train_config.do_cfg and unconditional_embeds is not None: unconditional_embeds = self.adapter.condition_encoded_embeds( tensors_0_1=clip_images, prompt_embeds=unconditional_embeds, is_training=True, has_been_preprocessed=True, is_unconditional=True, quad_count=quad_count ) if self.adapter and isinstance(self.adapter, CustomAdapter) and batch.extra_values is not None: self.adapter.add_extra_values(batch.extra_values.detach()) if self.train_config.do_cfg: self.adapter.add_extra_values(torch.zeros_like(batch.extra_values.detach()), is_unconditional=True) if has_adapter_img: if (self.adapter and isinstance(self.adapter, ControlNetModel)) or ( self.assistant_adapter and isinstance(self.assistant_adapter, ControlNetModel)): if self.train_config.do_cfg: raise ValueError("ControlNetModel is not supported with CFG") with torch.set_grad_enabled(self.adapter is not None): adapter: ControlNetModel = self.assistant_adapter if self.assistant_adapter is not None else self.adapter adapter_multiplier = get_adapter_multiplier() with self.timer('encode_adapter'): # add_text_embeds is pooled_prompt_embeds for sdxl added_cond_kwargs = {} if self.sd.is_xl: added_cond_kwargs["text_embeds"] = conditional_embeds.pooled_embeds added_cond_kwargs['time_ids'] = self.sd.get_time_ids_from_latents(noisy_latents) down_block_res_samples, mid_block_res_sample = adapter( noisy_latents, timesteps, encoder_hidden_states=conditional_embeds.text_embeds, controlnet_cond=adapter_images, conditioning_scale=1.0, guess_mode=False, added_cond_kwargs=added_cond_kwargs, return_dict=False, ) pred_kwargs['down_block_additional_residuals'] = down_block_res_samples pred_kwargs['mid_block_additional_residual'] = mid_block_res_sample if self.train_config.do_guidance_loss and isinstance(self.train_config.guidance_loss_target, list): batch_size = noisy_latents.shape[0] # update the guidance value, random float between guidance_loss_target[0] and guidance_loss_target[1] self._guidance_loss_target_batch = [ random.uniform( self.train_config.guidance_loss_target[0], self.train_config.guidance_loss_target[1] ) for _ in range(batch_size) ] self.before_unet_predict() if unconditional_embeds is not None: unconditional_embeds = unconditional_embeds.to(self.device_torch, dtype=dtype).detach() with self.timer('condition_noisy_latents'): # do it for the model noisy_latents = self.sd.condition_noisy_latents(noisy_latents, batch) if self.adapter and isinstance(self.adapter, CustomAdapter): noisy_latents = self.adapter.condition_noisy_latents(noisy_latents, batch) if self.train_config.timestep_type == 'next_sample': with self.timer('next_sample_step'): with torch.no_grad(): stepped_timestep_indicies = [self.sd.noise_scheduler.index_for_timestep(t) + 1 for t in timesteps] stepped_timesteps = [self.sd.noise_scheduler.timesteps[x] for x in stepped_timestep_indicies] stepped_timesteps = torch.stack(stepped_timesteps, dim=0) # do a sample at the current timestep and step it, then determine new noise next_sample_pred = self.predict_noise( noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype), timesteps=timesteps, conditional_embeds=conditional_embeds.to(self.device_torch, dtype=dtype), unconditional_embeds=unconditional_embeds, batch=batch, **pred_kwargs ) stepped_latents = self.sd.step_scheduler( next_sample_pred, noisy_latents, timesteps, self.sd.noise_scheduler ) # stepped latents is our new noisy latents. Now we need to determine noise in the current sample noisy_latents = stepped_latents original_samples = batch.latents.to(self.device_torch, dtype=dtype) # todo calc next timestep, for now this may work as it t_01 = (stepped_timesteps / 1000).to(original_samples.device) if len(stepped_latents.shape) == 4: t_01 = t_01.view(-1, 1, 1, 1) elif len(stepped_latents.shape) == 5: t_01 = t_01.view(-1, 1, 1, 1, 1) else: raise ValueError("Unknown stepped latents shape", stepped_latents.shape) next_sample_noise = (stepped_latents - (1.0 - t_01) * original_samples) / t_01 noise = next_sample_noise timesteps = stepped_timesteps # do a prior pred if we have an unconditional image, we will swap out the giadance later if batch.unconditional_latents is not None or self.do_guided_loss: # do guided loss loss = self.get_guided_loss( noisy_latents=noisy_latents, conditional_embeds=conditional_embeds, match_adapter_assist=match_adapter_assist, network_weight_list=network_weight_list, timesteps=timesteps, pred_kwargs=pred_kwargs, batch=batch, noise=noise, unconditional_embeds=unconditional_embeds, mask_multiplier=mask_multiplier, prior_pred=prior_pred, ) elif self.train_config.loss_type == 'mean_flow': loss = self.get_mean_flow_loss( noisy_latents=noisy_latents, conditional_embeds=conditional_embeds, match_adapter_assist=match_adapter_assist, network_weight_list=network_weight_list, timesteps=timesteps, pred_kwargs=pred_kwargs, batch=batch, noise=noise, unconditional_embeds=unconditional_embeds, prior_pred=prior_pred, ) else: with self.timer('predict_unet'): noise_pred = self.predict_noise( noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype), timesteps=timesteps, conditional_embeds=conditional_embeds.to(self.device_torch, dtype=dtype), unconditional_embeds=unconditional_embeds, batch=batch, is_primary_pred=True, **pred_kwargs ) self.after_unet_predict() with self.timer('calculate_loss'): noise = noise.to(self.device_torch, dtype=dtype).detach() prior_to_calculate_loss = prior_pred # if we are doing diff_output_preservation and not noing inverted masked prior # then we need to send none here so it will not target the prior if self.train_config.diff_output_preservation and not do_inverted_masked_prior: prior_to_calculate_loss = None loss = self.calculate_loss( noise_pred=noise_pred, noise=noise, noisy_latents=noisy_latents, timesteps=timesteps, batch=batch, mask_multiplier=mask_multiplier, prior_pred=prior_to_calculate_loss, ) if self.train_config.diff_output_preservation: # send the loss backwards otherwise checkpointing will fail self.accelerator.backward(loss) normal_loss = loss.detach() # dont send backward again dop_embeds = self.diff_output_preservation_embeds.expand_to_batch(noisy_latents.shape[0]) dop_pred = self.predict_noise( noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype), timesteps=timesteps, conditional_embeds=dop_embeds.to(self.device_torch, dtype=dtype), unconditional_embeds=unconditional_embeds, batch=batch, **pred_kwargs ) dop_loss = torch.nn.functional.mse_loss(dop_pred, prior_pred) * self.train_config.diff_output_preservation_multiplier self.accelerator.backward(dop_loss) loss = normal_loss + dop_loss loss = loss.clone().detach() # require grad again so the backward wont fail loss.requires_grad_(True) # check if nan if torch.isnan(loss): print_acc("loss is nan") loss = torch.zeros_like(loss).requires_grad_(True) with self.timer('backward'): # todo we have multiplier seperated. works for now as res are not in same batch, but need to change loss = loss * loss_multiplier.mean() # IMPORTANT if gradient checkpointing do not leave with network when doing backward # it will destroy the gradients. This is because the network is a context manager # and will change the multipliers back to 0.0 when exiting. They will be # 0.0 for the backward pass and the gradients will be 0.0 # I spent weeks on fighting this. DON'T DO IT # with fsdp_overlap_step_with_backward(): # if self.is_bfloat: # loss.backward() # else: self.accelerator.backward(loss) return loss.detach() # flush() def hook_train_loop(self, batch: Union[DataLoaderBatchDTO, List[DataLoaderBatchDTO]]): if isinstance(batch, list): batch_list = batch else: batch_list = [batch] total_loss = None self.optimizer.zero_grad() for batch in batch_list: if self.sd.is_multistage: # handle multistage switching if self.steps_this_boundary >= self.train_config.switch_boundary_every or self.current_boundary_index not in self.sd.trainable_multistage_boundaries: # iterate to make sure we only train trainable_multistage_boundaries while True: self.steps_this_boundary = 0 self.current_boundary_index += 1 if self.current_boundary_index >= len(self.sd.multistage_boundaries): self.current_boundary_index = 0 if self.current_boundary_index in self.sd.trainable_multistage_boundaries: # if this boundary is trainable, we can stop looking break loss = self.train_single_accumulation(batch) self.steps_this_boundary += 1 if total_loss is None: total_loss = loss else: total_loss += loss if len(batch_list) > 1 and self.model_config.low_vram: torch.cuda.empty_cache() if not self.is_grad_accumulation_step: # fix this for multi params if self.train_config.optimizer != 'adafactor': if isinstance(self.params[0], dict): for i in range(len(self.params)): self.accelerator.clip_grad_norm_(self.params[i]['params'], self.train_config.max_grad_norm) else: self.accelerator.clip_grad_norm_(self.params, self.train_config.max_grad_norm) # only step if we are not accumulating with self.timer('optimizer_step'): self.optimizer.step() self.optimizer.zero_grad(set_to_none=True) if self.adapter and isinstance(self.adapter, CustomAdapter): self.adapter.post_weight_update() if self.ema is not None: with self.timer('ema_update'): self.ema.update() else: # gradient accumulation. Just a place for breakpoint pass # TODO Should we only step scheduler on grad step? If so, need to recalculate last step with self.timer('scheduler_step'): self.lr_scheduler.step() if self.embedding is not None: with self.timer('restore_embeddings'): # Let's make sure we don't update any embedding weights besides the newly added token self.embedding.restore_embeddings() if self.adapter is not None and isinstance(self.adapter, ClipVisionAdapter): with self.timer('restore_adapter'): # Let's make sure we don't update any embedding weights besides the newly added token self.adapter.restore_embeddings() loss_dict = OrderedDict( {'loss': (total_loss / len(batch_list)).item()} ) self.end_of_training_loop() return loss_dict