up

app-fast.py

@@ -50,6 +50,7 @@ text_encoder = AutoModelForCausalLM.from_pretrained(
     output_attentions=True,
     low_cpu_mem_usage=True,
     quantization_config=quant_config,
+    device_map="auto",
     torch_dtype=torch.bfloat16,
 )
 
@@ -74,7 +75,7 @@ pipe = HiDreamImagePipeline.from_pretrained(
     scheduler=scheduler,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
-    device_map="auto",
+    device_map="balanced",
     torch_dtype=torch.bfloat16,
 ).to(device)