up

app-fast.py

@@ -2,19 +2,16 @@ import gradio as gr
 import PIL
 import spaces
 import torch
-from diffusers import TorchAoConfig as DiffusersTorchAoConfig
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
 from hi_diffusers.schedulers.flash_flow_match import (
     FlashFlowMatchEulerDiscreteScheduler,
 )
-from torchao.quantization import Int4WeightOnlyConfig
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
 )
-from transformers import (
-    TorchAoConfig as TransformersTorchAoConfig,
-)
+from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 
 # Constants
 MODEL_PREFIX: str = "HiDream-ai"
@@ -41,8 +38,10 @@ RESOLUTION_OPTIONS: list[str] = [
 
 device = torch.device("cuda")
 
-quant_config = Int4WeightOnlyConfig(group_size=128)
-quantization_config = TransformersTorchAoConfig(quant_type=quant_config)
+quant_config = TransformersBitsAndBytesConfig(
+    load_in_8bit=True,
+)
+
 
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
 text_encoder = AutoModelForCausalLM.from_pretrained(
@@ -50,18 +49,18 @@ text_encoder = AutoModelForCausalLM.from_pretrained(
     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
+    quantization_config=quant_config,
+    torch_dtype=torch.float16,
 ).to(device)  # Move model to the correct device after loading
 
-quantization_config = DiffusersTorchAoConfig("int8wo")
+quant_config = DiffusersBitsAndBytesConfig(
+    load_in_8bit=True,
+)
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
     MODEL_PATH,
     subfolder="transformer",
-    quantization_config=quantization_config,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
+    quantization_config=quant_config,
+    torch_dtype=torch.float16,
 ).to(device)
 
 scheduler = MODEL_CONFIGS["scheduler"](
@@ -75,7 +74,7 @@ pipe = HiDreamImagePipeline.from_pretrained(
     scheduler=scheduler,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float16,
 ).to(device)
 
 pipe.transformer = transformer

pyproject.toml

@@ -9,7 +9,6 @@ dependencies = [
     "diffusers>=0.32.1",
     "einops>=0.7.0",
     "torch>=2.5.1",
-    "torchao>=0.10.0",
     "torchvision>=0.20.1",
     "transformers>=4.47.1",
 ]

requirements.txt

@@ -10,4 +10,3 @@ einops
 gradio
 spaces
 sentencepiece
-torchao