load_quantized_instruct_nf4.py

"""
Quick loader for NF4 quantized HunyuanImage-3.0-Instruct model.
Generated automatically by hunyuan_quantize_instruct_nf4.py
"""

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

def load_quantized_instruct_nf4(model_path="H:\Testing\HunyuanImage-3.0-Instruct-NF4"):
    """Load the NF4 quantized HunyuanImage-3.0-Instruct model."""
    
    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        quantization_config=quant_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa",
    )
    
    # Load tokenizer
    model.load_tokenizer(model_path)
    
    return model

if __name__ == "__main__":
    print("Loading NF4 quantized Instruct model...")
    model = load_quantized_instruct_nf4()
    print("Model loaded successfully!")
    print(f"Device map: {model.hf_device_map}")
    
    if torch.cuda.is_available():
        print(f"GPU memory allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
        print(f"GPU memory reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")