Correct multi gpu dreambooth (#3673)

patrickvonplaten · web-flow · commit 262d539a8a8f · 2023-06-05T11:03:11.000+01:00
Correct multi gpu
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
@@ -1211,7 +1211,7 @@ def compute_text_embeddings(prompt):
                         text_encoder_use_attention_mask=args.text_encoder_use_attention_mask,
                     )
 
-                if unet.config.in_channels == channels * 2:
+                if accelerator.unwrap_model(unet).config.in_channels == channels * 2:
                     noisy_model_input = torch.cat([noisy_model_input, noisy_model_input], dim=1)
 
                 if args.class_labels_conditioning == "timesteps":
diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py
@@ -1156,7 +1156,7 @@ def compute_text_embeddings(prompt):
                         text_encoder_use_attention_mask=args.text_encoder_use_attention_mask,
                     )
 
-                if unet.config.in_channels == channels * 2:
+                if accelerator.unwrap_model(unet).config.in_channels == channels * 2:
                     noisy_model_input = torch.cat([noisy_model_input, noisy_model_input], dim=1)
 
                 if args.class_labels_conditioning == "timesteps":

Original file line number	Diff line number	Diff line change
`@@ -1211,7 +1211,7 @@ def compute_text_embeddings(prompt):`
`1211`	`1211`	`text_encoder_use_attention_mask=args.text_encoder_use_attention_mask,`
`1212`	`1212`	`)`
`1213`	`1213`
`1214`		`- if unet.config.in_channels == channels * 2:`
	`1214`	`+ if accelerator.unwrap_model(unet).config.in_channels == channels * 2:`
`1215`	`1215`	`noisy_model_input = torch.cat([noisy_model_input, noisy_model_input], dim=1)`
`1216`	`1216`
`1217`	`1217`	`if args.class_labels_conditioning == "timesteps":`
Original file line number	Diff line number	Diff line change
`@@ -1156,7 +1156,7 @@ def compute_text_embeddings(prompt):`
`1156`	`1156`	`text_encoder_use_attention_mask=args.text_encoder_use_attention_mask,`
`1157`	`1157`	`)`
`1158`	`1158`
`1159`		`- if unet.config.in_channels == channels * 2:`
	`1159`	`+ if accelerator.unwrap_model(unet).config.in_channels == channels * 2:`
`1160`	`1160`	`noisy_model_input = torch.cat([noisy_model_input, noisy_model_input], dim=1)`
`1161`	`1161`
`1162`	`1162`	`if args.class_labels_conditioning == "timesteps":`