import torch
from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig
from PIL import Image

# ===================== 1. 基础配置 =====================
# 你的本地路径
local_model_path = "/home/nashen/deepseek-ocr/DeepSeek-OCR-master/DeepSeek-OCR-vllm/model/"
device = "cuda" if torch.cuda.is_available() else "cpu"

# ===================== 2. 正确的分层量化配置 (BitsAndBytesConfig) =====================
# 核心修改：不再使用字典，而是创建一个 Config 对象
# 逻辑：默认全部 4bit 量化，然后通过 skip_modules 保护视觉层

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",       # 推荐 nf4
    bnb_4bit_use_double_quant=True,  # 双重量化省显存
    bnb_4bit_compute_dtype=torch.bfloat16,
    
    # 【重点】这里实现你的“分层”逻辑
    # 强制跳过视觉层(vision_model)和输出层(lm_head)，保持 bf16 原生精度
    llm_int8_skip_modules=[
        "vision_model",  # 视觉编码器
        "lm_head",       # 输出层
        "embed_tokens",  # 词嵌入
        "output"         # 部分架构的输出层
    ]
)

# ===================== 3. 加载模型 =====================
print(f"🚀 正在从本地加载模型: {local_model_path}")

try:
    tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)
    
    model = AutoModel.from_pretrained(
        local_model_path,
        trust_remote_code=True,
        # 核心修改：这里使用 quantization_config 参数，而不是 **kwargs 解包
        quantization_config=quantization_config, 
        device_map="auto"
    )
    
    # 开启 KV Cache 优化
    model.config.use_cache = True
    model.eval()

    print("\n✅ DeepSeek-OCR 本地模型加载成功！")
    print("✅ 混合精度策略已应用：")
    print("   - 视觉层/输出层：BF16 (原画质，不量化)")
    print("   - 语言层：NF4 (显存压缩)")
    
    # 测试一下显存占用
    print(f"✅ 模型加载完毕。你可以尝试运行推理了。")

except Exception as e:
    print(f"\n❌ 发生错误: {e}")