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

Fix issues found in review

Browse Source
This commit is contained in:
recris
2024-11-28 20:38:32 +00:00
parent 420a180d93
commit 740ec1d526
2 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
fine_tune.py
View File

@@ -411,7 +411,7 @@ def train(args):
loss = loss.mean() # mean over batch dimension
else:
loss = train_util.conditional_loss(
args, noise_pred.float(), target.float(), timesteps, "none", noise_scheduler
args, noise_pred.float(), target.float(), timesteps, "mean", noise_scheduler
)
accelerator.backward(loss)

8
library/train_util.py
View File

@@ -5829,8 +5829,8 @@ def save_sd_model_on_train_end_common(
def get_timesteps(min_timestep, max_timestep, b_size, device):
timesteps = torch.randint(min_timestep, max_timestep, (b_size,), device=device)
timesteps = timesteps.long()
timesteps = torch.randint(min_timestep, max_timestep, (b_size,), device="cpu")
timesteps = timesteps.long().to(device)
return timesteps
@@ -5875,8 +5875,8 @@ 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'):
raise NotImplementedError(f"Huber schedule 'snr' is not supported with the current model.")
if 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
result = (1 - args.huber_c) / (1 + sigmas) ** 2 + args.huber_c