VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add perturbation on noisy_model_input if needed

Browse Source
This commit is contained in:
rockerBOO
2025-03-19 00:37:27 -04:00
parent 7197266703
commit d93ad90a71
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
library/flux_train_utils.py
View File

@@ -410,20 +410,11 @@ def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None):
def get_noisy_model_input_and_timesteps( def get_noisy_model_input_and_timesteps(
args, noise_scheduler, latents: torch.Tensor, input_noise: torch.Tensor, device, dtype args, noise_scheduler, latents: torch.Tensor, noise: torch.Tensor, device, dtype
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
bsz, _, h, w = latents.shape bsz, _, h, w = latents.shape
sigmas = None sigmas = None
# Add noise to the latents according to the noise magnitude at each timestep
# (this is the forward diffusion process)
if args.ip_noise_gamma:
if args.ip_noise_gamma_random_strength:
noise = input_noise.detach().clone() + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents)
else:
noise = input_noise.detach().clone() + args.ip_noise_gamma * torch.randn_like(latents)
else:
noise = input_noise
if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid":
@@ -474,6 +465,15 @@ def get_noisy_model_input_and_timesteps(
sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype)
noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents
# Add noise to the latents according to the noise magnitude at each timestep
# (this is the forward diffusion process)
if args.ip_noise_gamma:
if args.ip_noise_gamma_random_strength:
xi = noise.detach().clone() + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents)
else:
xi = noise.detach().clone() + args.ip_noise_gamma * torch.randn_like(latents)
noisy_model_input += xi
return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas