up

app-fast.py

@@ -48,9 +48,9 @@ quantization_config = TransformersTorchAoConfig(
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
 text_encoder = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
-    torch_dtype=torch.bfloat16,
+    torch_dtype="auto",
     low_cpu_mem_usage=True,
-    device_map="auto",  # ✅ load directly onto GPU
+    device_map="auto",
     output_hidden_states=True,
     output_attentions=True,
     quantization_config=quantization_config,
@@ -60,9 +60,9 @@ quantization_config = DiffusersTorchAoConfig("int8wo")
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
     MODEL_PATH,
     subfolder="transformer",
-    device_map="auto",  # ✅ load directly onto GPU
+    device_map="auto",
     quantization_config=quantization_config,
-    torch_dtype=torch.bfloat16,
+    torch_dtype="auto",
 )
 
 scheduler = MODEL_CONFIGS["scheduler"](
@@ -76,8 +76,8 @@ pipe = HiDreamImagePipeline.from_pretrained(
     scheduler=scheduler,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
-    torch_dtype=torch.bfloat16,
-).to("cuda", torch.bfloat16)
+    torch_dtype="auto",
+)
 
 pipe.transformer = transformer