[+] Quantized inferencce

inference.py

@@ -4,32 +4,39 @@ from hi_diffusers import HiDreamImagePipeline
 from hi_diffusers import HiDreamImageTransformer2DModel
 from hi_diffusers.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
 from hi_diffusers.schedulers.flash_flow_match import FlashFlowMatchEulerDiscreteScheduler
-from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
+from transformers import LlamaForCausalLM, PreTrainedTokenizerFast, BitsAndBytesConfig
 parser = argparse.ArgumentParser()
 parser.add_argument("--model_type", type=str, default="dev")
 args = parser.parse_args()
 model_type = args.model_type
-MODEL_PREFIX = "HiDream-ai"
-LLAMA_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+
+
+MODEL_PREFIX = "azaneko"
+LLAMA_MODEL_NAME = "hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4"
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype="float16"
+)
 
 # Model configurations
 MODEL_CONFIGS = {
     "dev": {
-        "path": f"{MODEL_PREFIX}/HiDream-I1-Dev",
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Dev-nf4",
         "guidance_scale": 0.0,
         "num_inference_steps": 28,
         "shift": 6.0,
         "scheduler": FlashFlowMatchEulerDiscreteScheduler
     },
     "full": {
-        "path": f"{MODEL_PREFIX}/HiDream-I1-Full",
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Full-nf4",
         "guidance_scale": 5.0,
         "num_inference_steps": 50,
         "shift": 3.0,
         "scheduler": FlowUniPCMultistepScheduler
     },
     "fast": {
-        "path": f"{MODEL_PREFIX}/HiDream-I1-Fast",
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Fast-nf4",
         "guidance_scale": 0.0,
         "num_inference_steps": 16,
         "shift": 3.0,
@@ -37,69 +44,49 @@ MODEL_CONFIGS = {
     }
 }
 
-# Resolution options
-RESOLUTION_OPTIONS = [
-    "1024 × 1024 (Square)",
-    "768 × 1360 (Portrait)",
-    "1360 × 768 (Landscape)",
-    "880 × 1168 (Portrait)",
-    "1168 × 880 (Landscape)",
-    "1248 × 832 (Landscape)",
-    "832 × 1248 (Portrait)"
-]
+def log_vram(msg: str):
+    print(msg)
+    print(f"GPU memory usage: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
 
 # Load models
-def load_models(model_type):
+def load_models(model_type: str):
     config = MODEL_CONFIGS[model_type]
-    pretrained_model_name_or_path = config["path"]
-    scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=config["shift"], use_dynamic_shifting=False)
     
-    tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(
-        LLAMA_MODEL_NAME,
-        use_fast=False)
+    tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(LLAMA_MODEL_NAME)
+    log_vram("Tokenizer loaded!")
     
     text_encoder_4 = LlamaForCausalLM.from_pretrained(
         LLAMA_MODEL_NAME,
         output_hidden_states=True,
         output_attentions=True,
-        torch_dtype=torch.bfloat16).to("cuda")
+        return_dict_in_generate=True,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+    log_vram("Text encoder loaded!")
 
     transformer = HiDreamImageTransformer2DModel.from_pretrained(
-        pretrained_model_name_or_path, 
-        subfolder="transformer", 
-        torch_dtype=torch.bfloat16).to("cuda")
-
+        config["path"],
+        subfolder="transformer",
+        torch_dtype=torch.bfloat16
+    )
+    log_vram("Transformer loaded!")
+    
     pipe = HiDreamImagePipeline.from_pretrained(
-        pretrained_model_name_or_path, 
-        scheduler=scheduler,
+        config["path"],
+        scheduler=FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=config["shift"], use_dynamic_shifting=False),
         tokenizer_4=tokenizer_4,
         text_encoder_4=text_encoder_4,
-        torch_dtype=torch.bfloat16
-    ).to("cuda", torch.bfloat16)
+        torch_dtype=torch.bfloat16,
+    )
     pipe.transformer = transformer
+    log_vram("Pipeline loaded!")
+    pipe.enable_sequential_cpu_offload()
     
     return pipe, config
 
-# Parse resolution string to get height and width
-def parse_resolution(resolution_str):
-    if "1024 × 1024" in resolution_str:
-        return 1024, 1024
-    elif "768 × 1360" in resolution_str:
-        return 768, 1360
-    elif "1360 × 768" in resolution_str:
-        return 1360, 768
-    elif "880 × 1168" in resolution_str:
-        return 880, 1168
-    elif "1168 × 880" in resolution_str:
-        return 1168, 880
-    elif "1248 × 832" in resolution_str:
-        return 1248, 832
-    elif "832 × 1248" in resolution_str:
-        return 832, 1248
-    else:
-        return 1024, 1024  # Default fallback
-
 # Generate image function
+@torch.inference_mode()
 def generate_image(pipe, model_type, prompt, resolution, seed):
     # Get configuration for current model
     config = MODEL_CONFIGS[model_type]
@@ -107,8 +94,8 @@ def generate_image(pipe, model_type, prompt, resolution, seed):
     num_inference_steps = config["num_inference_steps"]
     
     # Parse resolution
-    height, width = parse_resolution(resolution)
-    
+    height, width = resolution
+ 
     # Handle seed
     if seed == -1:
         seed = torch.randint(0, 1000000, (1,)).item()
@@ -131,8 +118,9 @@ def generate_image(pipe, model_type, prompt, resolution, seed):
 print("Loading default model (full)...")
 pipe, _ = load_models(model_type)
 print("Model loaded successfully!")
-prompt = "A cat holding a sign that says \"Hi-Dreams.ai\"." 
-resolution = "1024 × 1024 (Square)"
+prompt = "A cat holding a sign that says \"I1 nf4\"." 
+# Possible values: 1024x1024, 768x1360, 1360x768, 880x1168, 1168x880, 1248x832, 832x1248
+resolution = (1024, 1024)
 seed = -1
 image, seed = generate_image(pipe, model_type, prompt, resolution, seed)
 image.save("output.png")