Update examples/dreambooth/train_dreambooth_lora_hidream.py

linoytsaban · hlky · web-flow · commit 677bab170edf · 2025-04-14T15:01:44.000+03:00
Co-authored-by: hlky &lt;hlky@hlky.ac&gt;
diff --git a/examples/dreambooth/train_dreambooth_lora_hidream.py b/examples/dreambooth/train_dreambooth_lora_hidream.py
@@ -1216,7 +1216,7 @@ def main(args):
         revision=args.revision,
         variant=args.variant,
     )
-    transformer = Lumina2Transformer2DModel.from_pretrained(
+    transformer = HiDreamImageTransformer2DModel.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant
     )
 

Original file line number	Diff line number	Diff line change
`@@ -1216,7 +1216,7 @@ def main(args):`
`1216`	`1216`	`revision=args.revision,`
`1217`	`1217`	`variant=args.variant,`
`1218`	`1218`	`)`
`1219`		`- transformer = Lumina2Transformer2DModel.from_pretrained(`
	`1219`	`+ transformer = HiDreamImageTransformer2DModel.from_pretrained(`
`1220`	`1220`	`args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant`
`1221`	`1221`	`)`
`1222`	`1222`