HiDream-ai-fast

Paused

blanchon commited on Apr 7

Commit

0343e87

1 Parent(s): f07518f

up

Files changed (1) hide show

app-fast.py CHANGED Viewed

@@ -75,7 +75,7 @@ pipe = HiDreamImagePipeline.from_pretrained(
     transformer=transformer,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
-    device_map="auto",
     torch_dtype=torch.bfloat16,
 ).to(device)

     transformer=transformer,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
+    device_map="balanced",
     torch_dtype=torch.bfloat16,
 ).to(device)