models/model_manager.py

import os
import hashlib
import logging
from typing import Optional, Dict, Any
from huggingface_hub import snapshot_download, HfApi, HfFolder
from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError

logger = logging.getLogger(__name__)

class ModelManager:
    def __init__(
        self,
        model_name: str = "qwen/qwen-image-2512",
        model_dir: str = "./models",
        revision: str = "main",
        cache_dir: Optional[str] = None
    ):
        self.model_name = model_name
        self.model_dir = model_dir
        self.revision = revision
        self.cache_dir = cache_dir or os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub")
        self.api = HfApi()
        
        # 确保模型目录存在
        os.makedirs(self.model_dir, exist_ok=True)
        
    def download_model(self) -> str:
        """
        下载模型到本地目录
        
        Returns:
            str: 下载的模型路径
        """
        try:
            logger.info(f"开始下载模型: {self.model_name} (版本: {self.revision})")
            
            # 使用snapshot_download下载模型，会自动处理缓存
            model_path = snapshot_download(
                repo_id=self.model_name,
                revision=self.revision,
                cache_dir=self.cache_dir,
                local_dir=self.model_dir,
                local_dir_use_symlinks=False,
                max_workers=4  # 限制下载线程数，避免资源占用过高
            )
            
            logger.info(f"模型下载完成: {model_path}")
            return model_path
            
        except RepositoryNotFoundError:
            logger.error(f"模型仓库不存在: {self.model_name}")
            raise
        except RevisionNotFoundError:
            logger.error(f"模型版本不存在: {self.revision}")
            raise
        except Exception as e:
            logger.error(f"模型下载失败: {str(e)}")
            raise
    
    def verify_model_integrity(self, model_path: str) -> bool:
        """
        验证模型文件完整性
        
        Args:
            model_path: 模型文件路径
            
        Returns:
            bool: 模型完整性验证结果
        """
        try:
            logger.info(f"验证模型完整性: {model_path}")
            
            # 检查必要的模型文件是否存在
            required_files = ["model.safetensors", "config.json", "tokenizer.json"]
            missing_files = []
            
            for file_name in required_files:
                file_path = os.path.join(model_path, file_name)
                if not os.path.exists(file_path):
                    missing_files.append(file_name)
            
            if missing_files:
                logger.error(f"缺少必要的模型文件: {missing_files}")
                return False
            
            logger.info("模型完整性验证通过")
            return True
            
        except Exception as e:
            logger.error(f"模型完整性验证失败: {str(e)}")
            return False
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息字典
        """
        try:
            logger.info(f"获取模型信息: {self.model_name}")
            
            model_info = self.api.model_info(
                repo_id=self.model_name,
                revision=self.revision
            )
            
            return {
                "model_name": model_info.id,
                "revision": model_info.sha,
                "created_at": model_info.created_at,
                "last_modified": model_info.last_modified,
                "size": model_info.siblings[0].size if model_info.siblings else 0,
                "tags": model_info.tags
            }
            
        except Exception as e:
            logger.error(f"获取模型信息失败: {str(e)}")
            raise
    
    def is_model_available(self) -> bool:
        """
        检查模型是否已存在于本地
        
        Returns:
            bool: 模型是否可用
        """
        try:
            # 检查模型目录是否存在且包含必要文件
            if not os.path.exists(self.model_dir):
                return False
            
            required_files = ["model.safetensors", "config.json"]
            for file_name in required_files:
                file_path = os.path.join(self.model_dir, file_name)
                if not os.path.exists(file_path):
                    return False
            
            return True
            
        except Exception as e:
            logger.error(f"检查模型可用性失败: {str(e)}")
            return False
    
    def calculate_file_hash(self, file_path: str, hash_algorithm: str = "sha256") -> str:
        """
        计算文件哈希值
        
        Args:
            file_path: 文件路径
            hash_algorithm: 哈希算法，默认为sha256
            
        Returns:
            str: 文件哈希值
        """
        try:
            hash_func = hashlib.new(hash_algorithm)
            with open(file_path, "rb") as f:
                # 分块读取文件，避免内存占用过高
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_func.update(chunk)
            
            return hash_func.hexdigest()
            
        except Exception as e:
            logger.error(f"计算文件哈希值失败: {str(e)}")
            raise
    
    def get_local_model_version(self) -> Optional[str]:
        """
        获取本地模型版本
        
        Returns:
            Optional[str]: 本地模型版本，如果不存在则返回None
        """
        try:
            version_file = os.path.join(self.model_dir, ".version")
            if os.path.exists(version_file):
                with open(version_file, "r") as f:
                    return f.read().strip()
            return None
        except Exception as e:
            logger.error(f"获取本地模型版本失败: {str(e)}")
            return None
    
    def save_model_version(self, version: str) -> None:
        """
        保存模型版本到本地
        
        Args:
            version: 模型版本
        """
        try:
            version_file = os.path.join(self.model_dir, ".version")
            with open(version_file, "w") as f:
                f.write(version)
            logger.info(f"模型版本已保存: {version}")
        except Exception as e:
            logger.error(f"保存模型版本失败: {str(e)}")
            raise
    
    def cleanup_old_models(self, keep_latest: int = 1) -> None:
        """
        清理旧模型版本
        
        Args:
            keep_latest: 保留的最新模型数量
        """
        try:
            # 这个功能可以根据需要扩展，目前简单实现
            logger.info(f"清理旧模型版本，保留最新 {keep_latest} 个版本")
            # 实际实现中可以遍历模型目录，根据版本号清理旧模型
        except Exception as e:
            logger.error(f"清理旧模型版本失败: {str(e)}")
            raise

if __name__ == "__main__":
    # 测试模型管理功能
    import logging
    logging.basicConfig(level=logging.INFO)
    
    model_manager = ModelManager()
    
    if not model_manager.is_model_available():
        model_path = model_manager.download_model()
        if model_manager.verify_model_integrity(model_path):
            model_info = model_manager.get_model_info()
            model_manager.save_model_version(model_info["revision"])
            print("模型下载和验证成功")
        else:
            print("模型验证失败")
    else:
        print("模型已存在")