up

app-fast.py

@@ -50,7 +50,6 @@ text_encoder = AutoModelForCausalLM.from_pretrained(
     output_attentions=True,
     low_cpu_mem_usage=True,
     quantization_config=quant_config,
-    device_map="auto",
     torch_dtype=torch.bfloat16,
 )
 
@@ -73,14 +72,13 @@ scheduler = MODEL_CONFIGS["scheduler"](
 pipe = HiDreamImagePipeline.from_pretrained(
     MODEL_PATH,
     scheduler=scheduler,
+    transformer=transformer,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
-    device_map="balanced",
+    device_map="auto",
     torch_dtype=torch.bfloat16,
 ).to(device)
 
-pipe.transformer = transformer
-
 
 @spaces.GPU(duration=120)
 def generate_image(