From 71c7fe216640a61bf01c016418c942c7ef012e9d Mon Sep 17 00:00:00 2001 From: Azalea <22280294+hykilpikonna@users.noreply.github.com> Date: Tue, 8 Apr 2025 21:25:09 +0000 Subject: [PATCH] [+] Quantized inferencce --- inference.py | 96 +++++++++++++++++++++++----------------------------- 1 file changed, 42 insertions(+), 54 deletions(-) diff --git a/inference.py b/inference.py index 1767dc9..757ecbc 100644 --- a/inference.py +++ b/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")