up

app-fast.py

@@ -42,18 +42,20 @@ RESOLUTION_OPTIONS: list[str] = [
 
 # Using AOBaseConfig instance (torchao >= 0.10.0)
 quant_config = Int4WeightOnlyConfig(group_size=128)
-quantization_config = TransformersTorchAoConfig(quant_type=quant_config)
+quantization_config = TransformersTorchAoConfig(
+    quant_type=quant_config, dtype=torch.bfloat16
+)
 
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
 text_encoder = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=True,
     device_map="auto",
     output_hidden_states=True,
     output_attentions=True,
     quantization_config=quantization_config,
-).to("cuda")
+).to("cuda", torch.bfloat16)
 
 quantization_config = DiffusersTorchAoConfig("int8wo")
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
@@ -61,7 +63,7 @@ transformer = HiDreamImageTransformer2DModel.from_pretrained(
     subfolder="transformer",
     quantization_config=quantization_config,
     torch_dtype=torch.bfloat16,
-).to("cuda")
+).to("cuda", dtype=torch.float16)
 
 scheduler = MODEL_CONFIGS["scheduler"](
     num_train_timesteps=1000,