From 8d5a183cc5cf8b00aa27c1cb2013fef46e34e3e4 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 20 Mar 2025 15:40:26 -0400 Subject: [PATCH] Fix applying image size to post_process_loss --- flux_train_network.py | 20 +- library/custom_train_functions.py | 85 +++++--- library/sd3_train_utils.py | 322 ++++++++++++++++++++++++++---- library/train_util.py | 27 ++- sd3_train_network.py | 2 +- train_network.py | 6 +- 6 files changed, 383 insertions(+), 79 deletions(-) diff --git a/flux_train_network.py b/flux_train_network.py index 712d0bc8..b875f678 100644 --- a/flux_train_network.py +++ b/flux_train_network.py @@ -22,7 +22,7 @@ from library import ( train_util, ) from library.custom_train_functions import ( - prepare_scheduler_for_custom_training, + prepare_scheduler_for_custom_training_flux, apply_snr_weight, scale_v_prediction_loss_like_noise_prediction, add_v_prediction_like_loss, @@ -331,9 +331,9 @@ class FluxNetworkTrainer(train_network.NetworkTrainer): """ def get_noise_scheduler(self, args: argparse.Namespace, device: torch.device) -> Any: - noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.discrete_flow_shift) + noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.discrete_flow_shift, use_dynamic_shifting=args.timestep_sampling == "flux_shift") self.noise_scheduler_copy = copy.deepcopy(noise_scheduler) - prepare_scheduler_for_custom_training(noise_scheduler, device) + prepare_scheduler_for_custom_training_flux(noise_scheduler, device) return noise_scheduler def encode_images_to_latents(self, args, vae, images): @@ -458,15 +458,19 @@ class FluxNetworkTrainer(train_network.NetworkTrainer): return model_pred, target, timesteps, weighting - def post_process_loss(self, loss: torch.Tensor, args, timesteps, noise_scheduler) -> torch.FloatTensor: + def post_process_loss(self, loss: torch.Tensor, args, timesteps, noise_scheduler, latents: Optional[torch.Tensor]) -> torch.FloatTensor: + image_size = None + if latents is not None: + image_size = tuple(latents.shape[-2:]) + if args.min_snr_gamma: - loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization) + loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization, image_size) if args.scale_v_pred_loss_like_noise_pred: - loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler) + loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler, image_size) if args.v_pred_like_loss: - loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss) + loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss, image_size) if args.debiased_estimation_loss: - loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization) + loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization, image_size) return loss def get_sai_model_spec(self, args): diff --git a/library/custom_train_functions.py b/library/custom_train_functions.py index 2d683693..2a657a9f 100644 --- a/library/custom_train_functions.py +++ b/library/custom_train_functions.py @@ -18,6 +18,9 @@ def prepare_scheduler_for_custom_training(noise_scheduler, device): if hasattr(noise_scheduler, "all_snr"): return + if hasattr(noise_scheduler.config, "use_dynamic_shifting") and noise_scheduler.config.use_dynamic_shifting is True: + return + alphas_cumprod = train_util.get_alphas_cumprod(noise_scheduler) sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) @@ -27,6 +30,22 @@ def prepare_scheduler_for_custom_training(noise_scheduler, device): noise_scheduler.all_snr = all_snr.to(device) +def prepare_scheduler_for_custom_training_flux(noise_scheduler, device): + if hasattr(noise_scheduler, "all_snr"): + return + + if hasattr(noise_scheduler.config, "use_dynamic_shifting") and noise_scheduler.config.use_dynamic_shifting is True: + return + + alphas_cumprod = train_util.get_alphas_cumprod(noise_scheduler) + if alphas_cumprod is None: + return + + sigma = 1.0 - alphas_cumprod + all_snr = (alphas_cumprod / sigma) + + noise_scheduler.all_snr = all_snr.to(device) + def fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler): # fix beta: zero terminal SNR @@ -66,9 +85,14 @@ def fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler): noise_scheduler.alphas_cumprod = alphas_cumprod -def apply_snr_weight(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, gamma: Number, v_prediction=False): - timesteps_indices = train_util.timesteps_to_indices(timesteps, len(noise_scheduler.all_snr)) - snr = torch.stack([noise_scheduler.all_snr[t] for t in timesteps_indices]) +def apply_snr_weight(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, gamma: Number, v_prediction=False, image_size=None): + # Get the appropriate SNR values based on timesteps and potentially image size + if hasattr(noise_scheduler, "get_snr_for_timestep"): + snr = noise_scheduler.get_snr_for_timestep(timesteps, image_size) + else: + timesteps_indices = train_util.timesteps_to_indices(timesteps, len(noise_scheduler.all_snr)) + snr = torch.stack([noise_scheduler.all_snr[t] for t in timesteps_indices]) + min_snr_gamma = torch.minimum(snr, torch.full_like(snr, gamma)) if v_prediction: snr_weight = torch.div(min_snr_gamma, snr + 1).float().to(loss.device) @@ -78,14 +102,19 @@ def apply_snr_weight(loss: torch.Tensor, timesteps: torch.IntTensor, noise_sched return loss -def scale_v_prediction_loss_like_noise_prediction(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler): - scale = get_snr_scale(timesteps, noise_scheduler) +def scale_v_prediction_loss_like_noise_prediction(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, image_size=None): + scale = get_snr_scale(timesteps, noise_scheduler, image_size) loss = loss * scale return loss -def get_snr_scale(timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler): - timesteps_indices = train_util.timesteps_to_indices(timesteps, len(noise_scheduler.all_snr)) - snr_t = torch.stack([noise_scheduler.all_snr[t] for t in timesteps_indices]) # batch_size +def get_snr_scale(timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, image_size=None): + # Get SNR values with image_size consideration + if hasattr(noise_scheduler, "get_snr_for_timestep"): + snr_t = noise_scheduler.get_snr_for_timestep(timesteps, image_size) + else: + timesteps_indices = train_util.timesteps_to_indices(timesteps, len(noise_scheduler.all_snr)) + snr_t = torch.stack([noise_scheduler.all_snr[t] for t in timesteps_indices]) + snr_t = torch.minimum(snr_t, torch.ones_like(snr_t) * 1000) # if timestep is 0, snr_t is inf, so limit it to 1000 scale = snr_t / (snr_t + 1) # # show debug info @@ -93,27 +122,37 @@ def get_snr_scale(timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler): return scale -def add_v_prediction_like_loss(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, v_pred_like_loss: torch.Tensor): - scale = get_snr_scale(timesteps, noise_scheduler) +def add_v_prediction_like_loss(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, v_pred_like_loss: torch.Tensor, image_size=None): + scale = get_snr_scale(timesteps, noise_scheduler, image_size) # logger.info(f"add v-prediction like loss: {v_pred_like_loss}, scale: {scale}, loss: {loss}, time: {timesteps}") loss = loss + loss / scale * v_pred_like_loss return loss -def apply_debiased_estimation(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, v_prediction=False): - if not hasattr(noise_scheduler, "all_snr"): - return loss +def apply_debiased_estimation(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, v_prediction=False, image_size=None): + # Check if we have SNR values available + if not (hasattr(noise_scheduler, "all_snr") or hasattr(noise_scheduler, "get_snr_for_timestep")): + return loss + + # Get SNR values with image_size consideration + if hasattr(noise_scheduler, "get_snr_for_timestep"): + snr_t = noise_scheduler.get_snr_for_timestep(timesteps, image_size) + else: + timesteps_indices = train_util.timesteps_to_indices(timesteps, len(noise_scheduler.all_snr)) + snr_t = torch.stack([noise_scheduler.all_snr[t] for t in timesteps_indices]) + + # Cap the SNR to avoid numerical issues + snr_t = torch.minimum(snr_t, torch.ones_like(snr_t) * 1000) + + # Apply weighting based on prediction type + if v_prediction: + weight = 1 / (snr_t + 1) + else: + weight = 1 / torch.sqrt(snr_t) + + loss = weight * loss + return loss - timesteps_indices = train_util.timesteps_to_indices(timesteps, len(noise_scheduler.all_snr)) - - snr_t = torch.stack([noise_scheduler.all_snr[t] for t in timesteps_indices]) # batch_size - snr_t = torch.minimum(snr_t, torch.ones_like(snr_t) * 1000) # if timestep is 0, snr_t is inf, so limit it to 1000 - if v_prediction: - weight = 1 / (snr_t + 1) - else: - weight = 1 / torch.sqrt(snr_t) - loss = weight * loss - return loss # TODO train_utilと分散しているのでどちらかに寄せる diff --git a/library/sd3_train_utils.py b/library/sd3_train_utils.py index c4079884..1fb34c74 100644 --- a/library/sd3_train_utils.py +++ b/library/sd3_train_utils.py @@ -28,7 +28,7 @@ import logging logger = logging.getLogger(__name__) -from library import sd3_models, sd3_utils, strategy_base, train_util +from library import sd3_models, sd3_utils, strategy_base, train_util, flux_train_utils def save_models( @@ -598,16 +598,29 @@ def sample_image_inference( # region Diffusers +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from dataclasses import dataclass -from typing import Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import numpy as np import torch from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.schedulers.scheduling_utils import SchedulerMixin -from diffusers.utils.torch_utils import randn_tensor from diffusers.utils import BaseOutput +from diffusers.schedulers.scheduling_utils import SchedulerMixin @dataclass @@ -649,22 +662,49 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): self, num_train_timesteps: int = 1000, shift: float = 1.0, + use_dynamic_shifting=False, + base_shift: Optional[float] = 0.5, + max_shift: Optional[float] = 1.15, + base_image_seq_len: Optional[int] = 256, + max_image_seq_len: Optional[int] = 4096, + invert_sigmas: bool = False, + shift_terminal: Optional[float] = None, + use_karras_sigmas: Optional[bool] = False, + use_exponential_sigmas: Optional[bool] = False, + use_beta_sigmas: Optional[bool] = False, ): + if self.config.use_beta_sigmas and not is_scipy_available(): + raise ImportError("Make sure to install scipy if you want to use beta sigmas.") + if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + raise ValueError( + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." + ) timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) sigmas = timesteps / num_train_timesteps - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) self.timesteps = sigmas * num_train_timesteps self._step_index = None self._begin_index = None + self._shift = shift + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication self.sigma_min = self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() + @property + def shift(self): + """ + The value used for shifting. + """ + return self._shift + @property def step_index(self): """ @@ -690,6 +730,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): """ self._begin_index = begin_index + def set_shift(self, shift: float): + self._shift = shift + def scale_noise( self, sample: torch.FloatTensor, @@ -709,10 +752,31 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): `torch.FloatTensor`: A scaled input sample. """ - if self.step_index is None: - self._init_step_index(timestep) + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) - sigma = self.sigmas[self.step_index] sample = sigma * noise + (1.0 - sigma) * sample return sample @@ -720,7 +784,37 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): def _sigma_to_t(self, sigma): return sigma * self.config.num_train_timesteps - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: + r""" + Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config + value. + + Reference: + https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 + + Args: + t (`torch.Tensor`): + A tensor of timesteps to be stretched and shifted. + + Returns: + `torch.Tensor`: + A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. + """ + one_minus_z = 1 - t + scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) + stretched_t = 1 - (one_minus_z / scale_factor) + return stretched_t + + def set_timesteps( + self, + num_inference_steps: int = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[float] = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -730,18 +824,49 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") + + if sigmas is None: + timesteps = np.linspace( + self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + ) + + sigmas = timesteps / self.config.num_train_timesteps + else: + sigmas = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas) self.num_inference_steps = num_inference_steps - timesteps = np.linspace(self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps) + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + + if self.config.shift_terminal: + sigmas = self.stretch_shift_to_terminal(sigmas) + + if self.config.use_karras_sigmas: + sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + elif self.config.use_exponential_sigmas: + sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + elif self.config.use_beta_sigmas: + sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - sigmas = timesteps / self.config.num_train_timesteps - sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) - timesteps = sigmas * self.config.num_train_timesteps - self.timesteps = timesteps.to(device=device) - self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + if self.config.invert_sigmas: + sigmas = 1.0 - sigmas + timesteps = sigmas * self.config.num_train_timesteps + sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) + else: + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + + self.timesteps = timesteps.to(device=device) + self.sigmas = sigmas self._step_index = None self._begin_index = None @@ -807,7 +932,11 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): returned, otherwise a tuple is returned where the first element is the sample tensor. """ - if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor): + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): raise ValueError( ( "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" @@ -823,30 +952,10 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): sample = sample.to(torch.float32) sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] - gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 + prev_sample = sample + (sigma_next - sigma) * model_output - noise = randn_tensor(model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator) - - eps = noise * s_noise - sigma_hat = sigma * (gamma + 1) - - if gamma > 0: - sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - # NOTE: "original_sample" should not be an expected prediction_type but is left in for - # backwards compatibility - - # if self.config.prediction_type == "vector_field": - - denoised = sample - model_output * sigma - # 2. Convert to an ODE derivative - derivative = (sample - denoised) / sigma_hat - - dt = self.sigmas[self.step_index + 1] - sigma_hat - - prev_sample = sample + derivative * dt # Cast sample back to model compatible dtype prev_sample = prev_sample.to(model_output.dtype) @@ -858,9 +967,146 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + """Constructs the noise schedule of Karras et al. (2022).""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential + def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: + """Constructs an exponential noise schedule.""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta + def _convert_to_beta( + self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 + ) -> torch.Tensor: + """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.array( + [ + sigma_min + (ppf * (sigma_max - sigma_min)) + for ppf in [ + scipy.stats.beta.ppf(timestep, alpha, beta) + for timestep in 1 - np.linspace(0, 1, num_inference_steps) + ] + ] + ) + return sigmas + def __len__(self): return self.config.num_train_timesteps + def get_snr_for_timestep(self, timesteps: torch.IntTensor, image_size=None): + """ + Get the signal-to-noise ratio for given timesteps, with consideration for image size. + + Args: + timesteps: Batch of timesteps (already scaled values, timesteps = sigma * 1000.0) + image_size: Tuple of (height, width) or single int representing image dimensions + + Returns: + torch.Tensor: SNR values corresponding to the timesteps + """ + + if not hasattr(self, "all_snr"): + all_sigmas = self.sigmas + assert isinstance(all_sigmas, torch.Tensor), "FlowMatch scheduler sigmas are not tensors" + + # Apply appropriate shifting to sigmas + if image_size is not None and self.config.use_dynamic_shifting: + # Calculate mu based on image dimensions + if isinstance(image_size, (tuple, list)): + h, w = image_size + else: + h = w = image_size + + # Adjust for packed size + h = h // 2 + w = w // 2 + mu = flux_train_utils.get_lin_function(y1=0.5, y2=1.15)(h * w) + + # Apply time shifting to sigmas + shifted_all_sigmas = self.time_shift(mu, 1.0, all_sigmas) + elif not self.config.use_dynamic_shifting: + # already shifted + shifted_all_sigmas = all_sigmas + else: + shifted_all_sigmas = all_sigmas + + # Calculate SNR based on shifted sigma values + all_snr = ((1.0 - shifted_all_sigmas**2) / (shifted_all_sigmas**2)).to(timesteps.device) + + # If we are using dynamic shifting we can't store all the snr + if not self.config.use_dynamic_shifting: + self.all_snr = all_snr + else: + all_snr = self.all_snr + + + # Convert input timesteps to indices + # Assuming timesteps are in the range [0, 1000] and need to be mapped to indices + timestep_indices = (timesteps / 1000.0 * (len(all_snr.to(timesteps.device)) - 1)).long() + + # Get SNR values for the requested timesteps + requested_snr = all_snr[timestep_indices] + + return requested_snr + + def get_sigmas(noise_scheduler, timesteps, device, n_dim=4, dtype=torch.float32): sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype) diff --git a/library/train_util.py b/library/train_util.py index e92d4518..9390be6b 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -31,6 +31,7 @@ from packaging.version import Version import torch from library.device_utils import init_ipex, clean_memory_on_device +from library.sd3_train_utils import FlowMatchEulerDiscreteScheduler from library.strategy_base import LatentsCachingStrategy, TokenizeStrategy, TextEncoderOutputsCachingStrategy, TextEncodingStrategy init_ipex() @@ -60,7 +61,7 @@ from diffusers import ( KDPM2AncestralDiscreteScheduler, AutoencoderKL, ) -from library import custom_train_functions, sd3_utils +from library import custom_train_functions, sd3_utils, flux_train_utils from library.original_unet import UNet2DConditionModel from huggingface_hub import hf_hub_download import numpy as np @@ -5976,7 +5977,7 @@ def get_noise_noisy_latents_and_timesteps( return noise, noisy_latents, timesteps -def get_huber_threshold_if_needed(args, timesteps: torch.Tensor, noise_scheduler) -> Optional[torch.Tensor]: +def get_huber_threshold_if_needed(args, timesteps: torch.Tensor, latents: torch.Tensor, noise_scheduler) -> Optional[torch.Tensor]: if not (args.loss_type == "huber" or args.loss_type == "smooth_l1"): return None @@ -5985,12 +5986,23 @@ def get_huber_threshold_if_needed(args, timesteps: torch.Tensor, noise_scheduler alpha = -math.log(args.huber_c) / noise_scheduler.config.num_train_timesteps result = torch.exp(-alpha * timesteps) * args.huber_scale elif args.huber_schedule == "snr": - alphas_cumprod = get_alphas_cumprod(noise_scheduler) + if hasattr(noise_scheduler, "sigmas"): + # Need to adjust the timesteps based on the latent dimensions + if args.timestep_sampling == "flux_shift": + _, _, h, w = latents.shape + mu = flux_train_utils.get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) + alphas_cumprod = get_alphas_cumprod(noise_scheduler, mu) + else: + alphas_cumprod = get_alphas_cumprod(noise_scheduler) + else: + alphas_cumprod = get_alphas_cumprod(noise_scheduler) + if alphas_cumprod is None: raise NotImplementedError("Huber schedule 'snr' is not supported with the current model.") timesteps_indices = index_for_timesteps(timesteps, noise_scheduler) alphas_cumprod = torch.index_select(alphas_cumprod.to(timesteps.device), 0, timesteps_indices) sigmas = ((1.0 - alphas_cumprod) / alphas_cumprod) ** 0.5 + result = (1 - args.huber_c) / (1 + sigmas) ** 2 + args.huber_c result = result.to(timesteps.device) elif args.huber_schedule == "constant": @@ -6039,7 +6051,7 @@ def timesteps_to_indices(timesteps: torch.Tensor, num_train_timesteps: int): return timesteps_indices -def get_alphas_cumprod(noise_scheduler) -> Optional[torch.Tensor]: +def get_alphas_cumprod(noise_scheduler, mu=None) -> Optional[torch.Tensor]: """ Get the cumulative product of the alpha values across the timesteps. @@ -6048,8 +6060,11 @@ def get_alphas_cumprod(noise_scheduler) -> Optional[torch.Tensor]: if hasattr(noise_scheduler, "alphas_cumprod"): alphas_cumprod = noise_scheduler.alphas_cumprod elif hasattr(noise_scheduler, "sigmas"): - # Since we don't have alphas_cumprod directly, we can derive it from sigmas - sigmas = noise_scheduler.sigmas + if noise_scheduler.config.use_dynamic_shifting is True: + sigmas = noise_scheduler.time_shift(mu, 1.0, noise_scheduler.sigmas) + else: + # Since we don't have alphas_cumprod directly, we can derive it from sigmas + sigmas = noise_scheduler.sigmas # In many diffusion models, sigma² = (1-α)/α where α is the cumulative product of alphas # So we can derive alphas_cumprod from sigmas diff --git a/sd3_train_network.py b/sd3_train_network.py index cdb7aa4e..01265733 100644 --- a/sd3_train_network.py +++ b/sd3_train_network.py @@ -391,7 +391,7 @@ class Sd3NetworkTrainer(train_network.NetworkTrainer): return model_pred, target, timesteps, weighting - def post_process_loss(self, loss, args, timesteps, noise_scheduler): + def post_process_loss(self, loss, args, timesteps, noise_scheduler, latents): return loss def get_sai_model_spec(self, args): diff --git a/train_network.py b/train_network.py index 2d279b3b..5a4d83f0 100644 --- a/train_network.py +++ b/train_network.py @@ -316,7 +316,7 @@ class NetworkTrainer: return noise_pred, target, timesteps, None - def post_process_loss(self, loss, args, timesteps: torch.IntTensor, noise_scheduler) -> torch.FloatTensor: + def post_process_loss(self, loss, args, timesteps: torch.IntTensor, noise_scheduler, latents: Optional[torch.Tensor]) -> torch.FloatTensor: if args.min_snr_gamma: loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization) if args.scale_v_pred_loss_like_noise_pred: @@ -442,7 +442,7 @@ class NetworkTrainer: is_train=is_train, ) - huber_c = train_util.get_huber_threshold_if_needed(args, timesteps, noise_scheduler) + huber_c = train_util.get_huber_threshold_if_needed(args, timesteps, latents, noise_scheduler) loss = train_util.conditional_loss(noise_pred.float(), target.float(), args.loss_type, "none", huber_c) if weighting is not None: loss = loss * weighting @@ -453,7 +453,7 @@ class NetworkTrainer: loss_weights = batch["loss_weights"] # 各sampleごとのweight loss = loss * loss_weights - loss = self.post_process_loss(loss, args, timesteps, noise_scheduler) + loss = self.post_process_loss(loss, args, timesteps, noise_scheduler, latents) return loss.mean()