StableDiffusion/WebUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Call weighted_forward during training

Browse Source
This commit is contained in:
Shondoit 2023-01-12 15:34:11 +01:00
parent 21642000b3
commit bc50936745
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
modules/hypernetworks/hypernetwork.py
View File

@ -640,13 +640,14 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
with devices.autocast(): with devices.autocast():
x = batch.latent_sample.to(devices.device, non_blocking=pin_memory) x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)
w = batch.weight.to(devices.device, non_blocking=pin_memory)
if tag_drop_out != 0 or shuffle_tags: if tag_drop_out != 0 or shuffle_tags:
shared.sd_model.cond_stage_model.to(devices.device) shared.sd_model.cond_stage_model.to(devices.device)
c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory) c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
shared.sd_model.cond_stage_model.to(devices.cpu) shared.sd_model.cond_stage_model.to(devices.cpu)
else: else:
c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory) c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory)
loss = shared.sd_model(x, c)[0] / gradient_step loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step
del x del x
del c del c

3
modules/textual_inversion/textual_inversion.py
View File

@ -480,6 +480,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
with devices.autocast(): with devices.autocast():
x = batch.latent_sample.to(devices.device, non_blocking=pin_memory) x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)
w = batch.weight.to(devices.device, non_blocking=pin_memory)
c = shared.sd_model.cond_stage_model(batch.cond_text) c = shared.sd_model.cond_stage_model(batch.cond_text)
if is_training_inpainting_model: if is_training_inpainting_model:
@ -490,7 +491,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
else: else:
cond = c cond = c
loss = shared.sd_model(x, cond)[0] / gradient_step loss = shared.sd_model.weighted_forward(x, cond, w)[0] / gradient_step
del x del x
_loss_step += loss.item() _loss_step += loss.item()