up

app-fast.py

@@ -2,15 +2,20 @@ import gradio as gr
 import PIL
 import spaces
 import torch
+from diffusers import TorchAoConfig
 from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
 from hi_diffusers.schedulers.flash_flow_match import (
     FlashFlowMatchEulerDiscreteScheduler,
 )
-from transformers import AutoTokenizer, LlamaForCausalLM
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AwqConfig,
+)
 
 # Constants
 MODEL_PREFIX: str = "HiDream-ai"
-LLAMA_MODEL_NAME: str = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+LLAMA_MODEL_NAME: str = "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4"
 MODEL_PATH = "HiDream-ai/HiDream-I1-Fast"
 MODEL_CONFIGS = {
     "guidance_scale": 0.0,
@@ -32,17 +37,30 @@ RESOLUTION_OPTIONS: list[str] = [
 ]
 
 
+quantization_config = AwqConfig(
+    bits=4,
+    fuse_max_seq_len=512,  # Note: Update this as per your use-case
+    do_fuse=True,
+)
+
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
-text_encoder = LlamaForCausalLM.from_pretrained(
+text_encoder = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    device_map="auto",
     output_hidden_states=True,
     output_attentions=True,
-    torch_dtype=torch.bfloat16,
+    quantization_config=quantization_config,
 ).to("cuda")
 
+quantization_config = TorchAoConfig("int8wo")
+
+quantization_config = TorchAoConfig("int8wo")
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
     MODEL_PATH,
     subfolder="transformer",
+    quantization_config=quantization_config,
     torch_dtype=torch.bfloat16,
 ).to("cuda")