michael/stable-diffusion-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

rework the code to not use the walrus operator because colab's 3.7 does not support it

Browse source
This commit is contained in:
AUTOMATIC 2022-11-05 17:09:42 +03:00
parent b8f2dfed3c
commit 62e3d71aa7
1 changed files with 5 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
modules/hypernetworks/hypernetwork.py
View file

@ -429,13 +429,16 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log
weights = hypernetwork.weights() weights = hypernetwork.weights()
for weight in weights: for weight in weights:
weight.requires_grad = True weight.requires_grad = True
# Here we use optimizer from saved HN, or we can specify as UI option. # Here we use optimizer from saved HN, or we can specify as UI option.
if (optimizer_name := hypernetwork.optimizer_name) in optimizer_dict: if hypernetwork.optimizer_name in optimizer_dict:
optimizer = optimizer_dict[hypernetwork.optimizer_name](params=weights, lr=scheduler.learn_rate) optimizer = optimizer_dict[hypernetwork.optimizer_name](params=weights, lr=scheduler.learn_rate)
optimizer_name = hypernetwork.optimizer_name
else: else:
print(f"Optimizer type {optimizer_name} is not defined!") print(f"Optimizer type {hypernetwork.optimizer_name} is not defined!")
optimizer = torch.optim.AdamW(params=weights, lr=scheduler.learn_rate) optimizer = torch.optim.AdamW(params=weights, lr=scheduler.learn_rate)
optimizer_name = 'AdamW' optimizer_name = 'AdamW'
if hypernetwork.optimizer_state_dict: # This line must be changed if Optimizer type can be different from saved optimizer. if hypernetwork.optimizer_state_dict: # This line must be changed if Optimizer type can be different from saved optimizer.
try: try:
optimizer.load_state_dict(hypernetwork.optimizer_state_dict) optimizer.load_state_dict(hypernetwork.optimizer_state_dict)