ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 11:11:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

More work on custom adapter

Browse Source
This commit is contained in:
Jaret Burkett
2024-01-16 17:41:26 -07:00
parent eebd3c8212
commit 655533d4c7
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
toolkit/train_tools.py
View File

@@ -788,7 +788,7 @@ def apply_snr_weight(
offset = 0
if noise_scheduler.timesteps[0] == 1000:
offset = 1
snr = torch.stack([all_snr[t - offset] for t in timesteps])
snr = torch.stack([all_snr[(t - offset).int()] for t in timesteps])
gamma_over_snr = torch.div(torch.ones_like(snr) * gamma, snr)
if fixed:
snr_weight = gamma_over_snr.float().to(loss.device) # directly using gamma over snr