mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-04-09 06:45:09 +00:00
Fix issues found in review (pt 2)
This commit is contained in:
@@ -5875,7 +5875,7 @@ def get_huber_threshold(args, timesteps: torch.Tensor, noise_scheduler) -> torch
|
||||
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":
|
||||
if not hasattr(noise_scheduler, "alphas_cumprod"):
|
||||
if noise_scheduler is None or not hasattr(noise_scheduler, "alphas_cumprod"):
|
||||
raise NotImplementedError("Huber schedule 'snr' is not supported with the current model.")
|
||||
alphas_cumprod = torch.index_select(noise_scheduler.alphas_cumprod, 0, timesteps.cpu())
|
||||
sigmas = ((1.0 - alphas_cumprod) / alphas_cumprod) ** 0.5
|
||||
|
||||
@@ -845,7 +845,7 @@ def train(args):
|
||||
# )
|
||||
# calculate loss
|
||||
loss = train_util.conditional_loss(
|
||||
args, model_pred.float(), target.float(), timesteps, "none", noise_scheduler
|
||||
args, model_pred.float(), target.float(), timesteps, "none", None
|
||||
)
|
||||
if args.masked_loss or ("alpha_masks" in batch and batch["alpha_masks"] is not None):
|
||||
loss = apply_masked_loss(loss, batch)
|
||||
|
||||
Reference in New Issue
Block a user