up

app-fast.py

@@ -40,7 +40,6 @@ RESOLUTION_OPTIONS: list[str] = [
 ]
 
 
-# Using AOBaseConfig instance (torchao >= 0.10.0)
 quant_config = Int4WeightOnlyConfig(group_size=128)
 quantization_config = TransformersTorchAoConfig(
     quant_type=quant_config, dtype=torch.bfloat16
@@ -51,19 +50,20 @@ text_encoder = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=True,
-    device_map=None,
+    device_map="auto",  # ✅ load directly onto GPU
     output_hidden_states=True,
     output_attentions=True,
     quantization_config=quantization_config,
-).to("cuda", torch.bfloat16)
+)
 
 quantization_config = DiffusersTorchAoConfig("int8wo")
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
     MODEL_PATH,
     subfolder="transformer",
+    device_map="auto",  # ✅ load directly onto GPU
     quantization_config=quantization_config,
     torch_dtype=torch.bfloat16,
-).to("cuda", dtype=torch.float16)
+)
 
 scheduler = MODEL_CONFIGS["scheduler"](
     num_train_timesteps=1000,