Update train_dreambooth.py

ethansmith2000 · web-flow · commit cbdfb09cd714 · 2024-06-10T21:04:53.000-04:00
fix variable name mse_loss_weights
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
@@ -1305,7 +1305,7 @@ def compute_text_embeddings(prompt):
                         # Velocity objective needs to be floored to an SNR weight of one.
                         snr = snr + 1             
                     
-                    base_weight = (
+                    mse_loss_weights = (
                         torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
                     )
 

Original file line number	Diff line number	Diff line change
`@@ -1305,7 +1305,7 @@ def compute_text_embeddings(prompt):`
`1305`	`1305`	`# Velocity objective needs to be floored to an SNR weight of one.`
`1306`	`1306`	`snr = snr + 1`
`1307`	`1307`
`1308`		`- base_weight = (`
	`1308`	`+ mse_loss_weights = (`
`1309`	`1309`	`torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr`
`1310`	`1310`	`)`
`1311`	`1311`