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

(ACTUAL) Min-SNR Weighting Strategy: Fixed SNR calculation to authors implementation

Browse Source
This commit is contained in:
AI-Casanova
2023-03-23 12:34:49 +00:00
parent a3c7d711e4
commit 518a18aeff
5 changed files with 17 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
library/custom_train_functions.py
View File

@@ -1,16 +1,20 @@
import torch
import argparse
import numpy as np
def apply_snr_weight(loss, latents, noisy_latents, gamma):
sigma = torch.sub(noisy_latents, latents) #find noise as applied by scheduler
zeros = torch.zeros_like(sigma)
alpha_mean_sq = torch.nn.functional.mse_loss(latents.float(), zeros.float(), reduction="none").mean([1, 2, 3]) #trick to get Mean Square/Second Moment
sigma_mean_sq = torch.nn.functional.mse_loss(sigma.float(), zeros.float(), reduction="none").mean([1, 2, 3]) #trick to get Mean Square/Second Moment
snr = torch.div(alpha_mean_sq,sigma_mean_sq) #Signal to Noise Ratio = ratio of Mean Squares
def apply_snr_weight(loss, timesteps, noise_scheduler, gamma):
alphas_cumprod = noise_scheduler.alphas_cumprod.cpu()
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - alphas_cumprod)
alpha = sqrt_alphas_cumprod
sigma = sqrt_one_minus_alphas_cumprod
all_snr = (alpha / sigma) ** 2
all_snr.to(loss.device)
snr = torch.stack([all_snr[t] for t in timesteps])
gamma_over_snr = torch.div(torch.ones_like(snr)*gamma,snr)
snr_weight = torch.minimum(gamma_over_snr,torch.ones_like(gamma_over_snr)).float() #from paper
snr_weight = torch.minimum(gamma_over_snr,torch.ones_like(gamma_over_snr)).float().to(loss.device) #from paper
loss = loss * snr_weight
#print(snr_weight)
return loss
def add_custom_train_arguments(parser: argparse.ArgumentParser):