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

Fix IP noise calculation

Browse Source
This commit is contained in:
rockerBOO
2025-03-19 13:50:04 -04:00
parent 1eddac26b0
commit 5d5a7d2acf
1 changed files with 6 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
library/flux_train_utils.py
View File

@@ -423,29 +423,24 @@ def get_noisy_model_input_and_timesteps(
else: else:
t = torch.rand((bsz,), device=device) t = torch.rand((bsz,), device=device)
sigmas = t.view(-1, 1, 1, 1)
timesteps = t * 1000.0 timesteps = t * 1000.0
t = t.view(-1, 1, 1, 1)
noisy_model_input = (1 - t) * latents + t * noise
elif args.timestep_sampling == "shift": elif args.timestep_sampling == "shift":
shift = args.discrete_flow_shift shift = args.discrete_flow_shift
logits_norm = torch.randn(bsz, device=device) logits_norm = torch.randn(bsz, device=device)
logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling
timesteps = logits_norm.sigmoid() timesteps = logits_norm.sigmoid()
timesteps = (timesteps * shift) / (1 + (shift - 1) * timesteps) timesteps = (timesteps * shift) / (1 + (shift - 1) * timesteps)
sigmas = timesteps.view(-1, 1, 1, 1)
t = timesteps.view(-1, 1, 1, 1)
timesteps = timesteps * 1000.0 timesteps = timesteps * 1000.0
noisy_model_input = (1 - t) * latents + t * noise
elif args.timestep_sampling == "flux_shift": elif args.timestep_sampling == "flux_shift":
logits_norm = torch.randn(bsz, device=device) logits_norm = torch.randn(bsz, device=device)
logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling
timesteps = logits_norm.sigmoid() timesteps = logits_norm.sigmoid()
mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2))
timesteps = time_shift(mu, 1.0, timesteps) timesteps = time_shift(mu, 1.0, timesteps)
sigmas = timesteps.view(-1, 1, 1, 1)
t = timesteps.view(-1, 1, 1, 1)
timesteps = timesteps * 1000.0 timesteps = timesteps * 1000.0
noisy_model_input = (1 - t) * latents + t * noise
else: else:
# Sample a random timestep for each image # Sample a random timestep for each image
# for weighting schemes where we sample timesteps non-uniformly # for weighting schemes where we sample timesteps non-uniformly
@@ -458,10 +453,7 @@ def get_noisy_model_input_and_timesteps(
) )
indices = (u * noise_scheduler.config.num_train_timesteps).long() indices = (u * noise_scheduler.config.num_train_timesteps).long()
timesteps = noise_scheduler.timesteps[indices].to(device=device) timesteps = noise_scheduler.timesteps[indices].to(device=device)
# Add noise according to flow matching.
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
# Add noise to the latents according to the noise magnitude at each timestep # Add noise to the latents according to the noise magnitude at each timestep
# (this is the forward diffusion process) # (this is the forward diffusion process)
@@ -471,7 +463,9 @@ def get_noisy_model_input_and_timesteps(
ip_noise_gamma = (torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma) ip_noise_gamma = (torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma)
else: else:
ip_noise_gamma = args.ip_noise_gamma ip_noise_gamma = args.ip_noise_gamma
noisy_model_input += ip_noise_gamma * xi noisy_model_input = sigmas * (noise + ip_noise_gamma * xi) + (1.0 - sigmas) * latents
else:
noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents
return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas