up

app-fast.py

@@ -39,31 +39,30 @@ RESOLUTION_OPTIONS: list[str] = [
     "832 x 1248 (Portrait)",
 ]
 
+device = torch.device("cuda")
 
 quant_config = Int4WeightOnlyConfig(group_size=128)
-quantization_config = TransformersTorchAoConfig(
-    quant_type=quant_config, dtype=torch.bfloat16
-)
+quantization_config = TransformersTorchAoConfig(quant_type=quant_config)
 
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
 text_encoder = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
-    torch_dtype="auto",
-    low_cpu_mem_usage=True,
-    device_map="auto",
     output_hidden_states=True,
     output_attentions=True,
+    low_cpu_mem_usage=True,
     quantization_config=quantization_config,
-)
+    torch_dtype=torch.bfloat16,  # Explicitly set dtype
+    device_map="auto",  # Still use auto, but ensure device consistency
+).to(device)  # Move model to the correct device after loading
 
 quantization_config = DiffusersTorchAoConfig("int8wo")
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
     MODEL_PATH,
     subfolder="transformer",
-    device_map="auto",
     quantization_config=quantization_config,
-    torch_dtype="auto",
-)
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+).to(device)
 
 scheduler = MODEL_CONFIGS["scheduler"](
     num_train_timesteps=1000,
@@ -76,8 +75,8 @@ pipe = HiDreamImagePipeline.from_pretrained(
     scheduler=scheduler,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
-    torch_dtype="auto",
-)
+    torch_dtype=torch.bfloat16,
+).to(device)
 
 pipe.transformer = transformer