| """
|
| 模型加载模块 (4-bit 量化)
|
| - 加载 Qwen2.5-VL 7B + bitsandbytes NF4 量化
|
| - 适配 Kaggle T4 15GB 显存
|
| - 单例模式，全局共享一个模型实例
|
| """
|
|
|
| import torch
|
| from typing import Optional, Tuple
|
|
|
| from PIL import Image
|
|
|
|
|
| _model: Optional[object] = None
|
| _processor: Optional[object] = None
|
|
|
|
|
| def load_model_and_processor(
|
| model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct",
|
| use_quantization: bool = True,
|
| device_map: str = "auto",
|
| ):
|
| """
|
| 加载 Qwen2.5-VL 模型 + 处理器
|
|
|
| Args:
|
| model_name: HuggingFace 模型名
|
| use_quantization: 是否使用 4-bit 量化
|
| device_map: 设备分配策略
|
|
|
| Returns:
|
| (model, processor)
|
| """
|
| global _model, _processor
|
|
|
| if _model is not None and _processor is not None:
|
| return _model, _processor
|
|
|
| from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
|
|
|
| print(f"[ModelLoader] 正在加载模型: {model_name}")
|
|
|
|
|
| if use_quantization:
|
| from transformers import BitsAndBytesConfig
|
|
|
| bnb_config = BitsAndBytesConfig(
|
| load_in_4bit=True,
|
| bnb_4bit_compute_dtype=torch.float16,
|
| bnb_4bit_use_double_quant=True,
|
| bnb_4bit_quant_type="nf4",
|
| )
|
|
|
| model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
| model_name,
|
| quantization_config=bnb_config,
|
| device_map=device_map,
|
| trust_remote_code=True,
|
| )
|
| print("[ModelLoader] 已启用 4-bit NF4 量化")
|
| else:
|
| model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
| model_name,
|
| torch_dtype=torch.float16,
|
| device_map=device_map,
|
| trust_remote_code=True,
|
| )
|
|
|
|
|
| processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
|
|
|
|
|
| if torch.cuda.is_available():
|
| allocated = torch.cuda.memory_allocated(0) / (1024**3)
|
| print(f"[ModelLoader] GPU 显存已占用: {allocated:.2f} GB")
|
|
|
| _model = model
|
| _processor = processor
|
| return model, processor
|
|
|
|
|
| def get_model():
|
| """获取已加载的模型（单例）"""
|
| global _model
|
| if _model is None:
|
| raise RuntimeError("模型尚未加载，请先调用 load_model_and_processor()")
|
| return _model
|
|
|
|
|
| def get_processor():
|
| """获取已加载的处理器（单例）"""
|
| global _processor
|
| if _processor is None:
|
| raise RuntimeError("处理器尚未加载，请先调用 load_model_and_processor()")
|
| return _processor
|
|
|
|
|
| def is_loaded() -> bool:
|
| """检查模型是否已加载"""
|
| return _model is not None and _processor is not None
|
|
|
