model_manager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import json
import shutil
import glob
import requests
import subprocess
from pathlib import Path
from typing import Optional, Callable, Dict, Any

# 尝试导入modelscope和pycsghub
try:
    from modelscope.hub.snapshot_download import snapshot_download
    has_modelscope = True
except ImportError:
    print("Warning: modelscope not installed, download functionality will be limited")
    has_modelscope = False

try:
    from pycsghub.upload_large_folder.main import upload_large_folder_internal, create_repo
    from pycsghub.csghub_api import CsgHubApi
    has_pycsghub = True
except ImportError:
    print("Warning: pycsghub not installed, upload functionality will be limited")
    has_pycsghub = False


class ModelManager:
    """模型管理器，用于下载和上传模型"""
    
    def __init__(self):
        """初始化模型管理器"""
        self.default_download_path = os.path.expanduser("~/models")
        self.csghub_config = {
            "base_url": "http://10.17.27.227:4997",
            "token": "f5dad38a9426410aa861155cd184f84a",
            "repo_type": "model",
            "revision": "main"
        }
        
        # 确保默认下载路径存在
        os.makedirs(self.default_download_path, exist_ok=True)
    
    def download_model(self, model_id: str, local_path: str = None, 
                      progress_callback: Optional[Callable] = None,
                      cancel_check: Optional[Callable] = None) -> str:
        """
        从ModelScope下载模型
        
        Args:
            model_id: 模型ID，格式为"组织/模型名"
            local_path: 本地保存路径，默认为~/models
            progress_callback: 进度回调函数，接收(progress, detail)参数
            cancel_check: 取消检查函数，返回True表示已取消
        
        Returns:
            str: 下载的模型路径
        
        Raises:
            Exception: 下载失败时抛出异常
        """
        if not model_id:
            raise ValueError("模型ID不能为空")
        
        # 设置本地路径
        if not local_path:
            local_path = self.default_download_path
        
        # 确保本地路径存在
        os.makedirs(local_path, exist_ok=True)
        
        # 模型保存的完整路径
        model_path = os.path.join(local_path, model_id)
        
        # 打印下载信息到终端
        print(f"\n{'='*50}")
        print(f"开始下载模型")
        print(f"模型ID: {model_id}")
        print(f"本地路径: {model_path}")
        print(f"{'='*50}")
        
        # 如果模型已存在，先删除
        if os.path.exists(model_path):
            print(f"模型已存在，正在删除: {model_path}")
            shutil.rmtree(model_path)
        
        # 调用回调函数
        if progress_callback:
            progress_callback(0, f"开始下载模型 {model_id}")
        
        try:
            if has_modelscope:
                # 使用modelscope下载
                print(f"使用modelscope下载模型: {model_id}")
                print(f"下载目标路径: {model_path}")
                
                # 检查是否已取消
                if cancel_check and cancel_check():
                    print(f"下载任务已取消: {model_id}")
                    if progress_callback:
                        progress_callback(-1, {"error": "下载已取消"})
                    raise Exception("下载已取消")
                
                # 注意：modelscope不直接支持进度回调，我们将在下载后计算文件数量
                # 使用进程池执行snapshot_download，以便可以强制终止
                import multiprocessing
                
                # 定义下载函数
                def download_func():
                    try:
                        return snapshot_download(
                            model_id=model_id,
                            cache_dir=local_path,
                            revision="master"
                        )
                    except Exception as e:
                        print(f"下载出错: {e}")
                        raise
                
                # 创建进程
                process = multiprocessing.Process(target=download_func)
                process.daemon = True
                process.start()
                
                # 定期检查是否已取消
                while process.is_alive():
                    if cancel_check and cancel_check():
                        print(f"下载任务已取消: {model_id}")
                        # 强制终止进程
                        process.terminate()
                        process.join(timeout=1)
                        if process.is_alive():
                            process.kill()
                        if progress_callback:
                            progress_callback(-1, {"error": "下载已取消"})
                        raise Exception("下载已取消")
                    time.sleep(0.1)
                
                # 检查进程是否正常退出
                if process.exitcode != 0:
                    raise Exception("下载进程异常退出")
                
                print(f"modelscope下载完成，正在处理文件...")
                
                # 下载完成后，计算文件数量并更新进度
                if os.path.exists(model_path):
                    # 获取文件列表
                    all_files = []
                    for root, dirs, files in os.walk(model_path):
                        # 检查是否已取消
                        if cancel_check and cancel_check():
                            print(f"下载任务已取消: {model_id}")
                            if progress_callback:
                                progress_callback(-1, {"error": "下载已取消"})
                            raise Exception("下载已取消")
                        for file in files:
                            all_files.append(os.path.join(root, file))
                    
                    file_count = len(all_files)
                    print(f"发现 {file_count} 个文件")
                    
                    # 按文件数量更新进度
                    for i, file_path in enumerate(all_files):
                        # 检查是否已取消
                        if cancel_check and cancel_check():
                            print(f"下载任务已取消: {model_id}")
                            if progress_callback:
                                progress_callback(-1, {"error": "下载已取消"})
                            raise Exception("下载已取消")
                        progress = int((i + 1) / file_count * 100)
                        rel_path = os.path.relpath(file_path, model_path)
                        file_size = os.path.getsize(file_path)
                        print(f"[{progress}%] 已下载: {rel_path} ({self.get_dir_size(file_path)})")
                        if progress_callback:
                            progress_callback(progress, {
                                "file_count": i + 1,
                                "total_files": file_count,
                                "current_file": rel_path,
                                "file_size": file_size
                            })
                        time.sleep(0.05)  # 减少延迟
            else:
                # 直接使用modelscope下载（不使用模拟模式）
                print(f"modelscope未安装，无法下载模型: {model_id}")
                raise Exception("modelscope未安装，无法下载模型")
            
            if progress_callback:
                progress_callback(100, {
                    "file_count": file_count,
                    "total_files": file_count,
                    "current_file": "完成",
                    "message": f"模型 {model_id} 下载完成"
                })
            
            # 打印下载完成信息
            print(f"\n{'='*50}")
            print(f"模型下载完成!")
            print(f"模型ID: {model_id}")
            print(f"下载路径: {model_path}")
            print(f"文件数量: {file_count}")
            print(f"{'='*50}")
            
            return model_path
            
        except Exception as e:
            error_msg = str(e)
            if progress_callback:
                progress_callback(-1, {"error": error_msg})
            
            # 打印错误信息
            print(f"\n{'='*50}")
            print(f"下载失败!")
            print(f"模型ID: {model_id}")
            print(f"错误信息: {error_msg}")
            print(f"{'='*50}")
            
            raise Exception(f"下载模型 {model_id} 失败: {error_msg}")
    
    def upload_model(self, local_path: str, repo_id: str, 
                    create_repo_flag: bool = True,
                    progress_callback: Optional[Callable] = None) -> Dict[str, Any]:
        """
        上传模型到CsgHub
        
        Args:
            local_path: 本地模型路径
            repo_id: 仓库ID
            create_repo_flag: 是否创建仓库
            progress_callback: 进度回调函数，接收(progress, detail)参数
        
        Returns:
            Dict[str, Any]: 上传结果
        
        Raises:
            Exception: 上传失败时抛出异常
        """
        if not local_path or not os.path.exists(local_path):
            raise ValueError(f"本地路径 {local_path} 不存在")
        
        if not repo_id:
            raise ValueError("仓库ID不能为空")
        
        # 调用回调函数
        if progress_callback:
            progress_callback(0, f"开始上传模型到仓库 {repo_id}")
        
        try:
            # 首先获取所有文件列表，用于计算进度
            all_files = []
            for root, dirs, files in os.walk(local_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    all_files.append(file_path)
            
            file_count = len(all_files)
            
            if file_count == 0:
                raise ValueError(f"本地路径 {local_path} 中没有文件")
            
            if has_pycsghub:
                # 使用pycsghub上传
                csg_api = CsgHubApi()
                use_full_repo_id = f"root/{repo_id}"
                
                # 创建仓库
                if create_repo_flag:
                    if progress_callback:
                        progress_callback(5, "正在创建仓库...")
                    
                    create_repo(
                        api=csg_api,
                        repo_id=use_full_repo_id,
                        repo_type=self.csghub_config["repo_type"],
                        revision=self.csghub_config["revision"],
                        endpoint=self.csghub_config["base_url"],
                        token=self.csghub_config["token"]
                    )
                
                # 上传模型
                if progress_callback:
                    progress_callback(10, f"准备上传 {file_count} 个文件...")
                
                # 创建一个自定义的进度回调函数
                def custom_upload_callback(current_file_index, current_file_path, total_files):
                    """自定义上传进度回调"""
                    progress = int((current_file_index + 1) / total_files * 90) + 10  # 10% - 100%
                    rel_path = os.path.relpath(current_file_path, local_path)
                    if progress_callback:
                        progress_callback(progress, f"上传中 {current_file_index + 1}/{total_files}: {rel_path}")
                
                # 执行上传 - 注意：pycsghub可能不直接支持文件级别的进度回调
                # 这里我们将在上传完成后模拟文件级别的进度
                upload_large_folder_internal(
                    repo_id=use_full_repo_id,
                    local_path=local_path,
                    repo_type=self.csghub_config["repo_type"],
                    revision=self.csghub_config["revision"],
                    endpoint=self.csghub_config["base_url"],
                    token=self.csghub_config["token"],
                    allow_patterns=None,
                    ignore_patterns=None,
                    num_workers=1,
                    print_report=False,
                    print_report_every=1,
                )
                
                # 上传完成后，模拟文件级别的进度更新
                for i, file_path in enumerate(all_files):
                    progress = int((i + 1) / file_count * 90) + 10  # 10% - 100%
                    rel_path = os.path.relpath(file_path, local_path)
                    if progress_callback:
                        progress_callback(progress, f"已上传 {i + 1}/{file_count}: {rel_path}")
                    time.sleep(0.05)  # 模拟处理延迟
            else:
                # 直接使用pycsghub上传（不使用模拟模式）
                print(f"pycsghub未安装，无法上传模型: {repo_id}")
                raise Exception("pycsghub未安装，无法上传模型")
            
            if progress_callback:
                progress_callback(100, f"模型上传完成，仓库ID: {repo_id}，共上传 {file_count} 个文件")
            
            return {
                "success": True,
                "repo_id": repo_id,
                "file_count": file_count,
                "message": f"模型上传成功，共上传 {file_count} 个文件"
            }
            
        except Exception as e:
            if progress_callback:
                progress_callback(-1, f"上传失败: {str(e)}")
            raise Exception(f"上传模型失败: {str(e)}")
    


    def list_models(self, local_path: str = None) -> list:
        """
        列出本地模型
        
        Args:
            local_path: 本地模型路径，默认为~/models
        
        Returns:
            list: 模型列表
        """
        if not local_path:
            local_path = self.default_download_path
        
        if not os.path.exists(local_path):
            print(f"Model path does not exist: {local_path}")
            return []
        
        models = []
        
        try:
            print(f"Listing models from: {local_path}")
            items = os.listdir(local_path)
            print(f"Found {len(items)} items in directory")
            
            # 遍历一级目录
            for item in items:
                item_path = os.path.join(local_path, item)
                print(f"Checking item: {item} (type: {'dir' if os.path.isdir(item_path) else 'file'})")
                
                if os.path.isdir(item_path):
                    # 检查一级目录下的二级子目录
                    try:
                        sub_items = os.listdir(item_path)
                        print(f"  Found {len(sub_items)} sub-items in {item}")
                        
                        for sub_item in sub_items:
                            sub_item_path = os.path.join(item_path, sub_item)
                            
                            if os.path.isdir(sub_item_path):
                                print(f"  Checking sub-directory: {sub_item}")
                                
                                # 检查是否有 README.md
                                has_readme = os.path.exists(os.path.join(sub_item_path, "README.md"))
                                
                                # 检查是否有 .safetensors 或 .bin 文件
                                has_safetensors_or_bin = False
                                try:
                                    for file in os.listdir(sub_item_path):
                                        if file.endswith('.safetensors') or file.endswith('.bin'):
                                            has_safetensors_or_bin = True
                                            break
                                except Exception as e:
                                    print(f"    Error checking files in {sub_item_path}: {e}")
                                    continue
                                
                                print(f"    - README.md: {has_readme}")
                                print(f"    - has .safetensors or .bin: {has_safetensors_or_bin}")
                                
                                # 判断是否为模型目录（必须有README.md和.safetensors/.bin文件）
                                if has_readme and has_safetensors_or_bin:
                                    # 获取模型信息，使用前端期望的字段名
                                    model_info = {
                                        "id": sub_item,  # 使用二级目录名作为id
                                        "path": sub_item_path,
                                        "size": self.get_dir_size(sub_item_path),
                                        "status": "downloaded",  # 默认状态
                                        "downloadTime": self.get_dir_creation_time(sub_item_path),
                                        "uploadTime": None,
                                        "upload_repo_id": None,
                                        "file_count": 0  # 计算文件数量
                                    }
                                    
                                    # 计算文件数量
                                    file_count = 0
                                    for root, dirs, files in os.walk(sub_item_path):
                                        file_count += len(files)
                                    model_info["file_count"] = file_count
                                    
                                    models.append(model_info)
                                    print(f"    + Added as model: {sub_item} (in {item}/{sub_item})")
                                else:
                                    print(f"    - Skipped (missing required files)")
                            else:
                                print(f"    - Skipped (not a directory): {sub_item}")
                                
                    except Exception as e:
                        print(f"  Error processing sub-directories in {item_path}: {e}")
                        continue
                else:
                    print(f"  - Skipped (not a directory): {item}")
            
            print(f"Total models found: {len(models)}")
            
        except Exception as e:
            print(f"Error listing models: {e}")
        
        return models
    
    def get_dir_size(self, path: str) -> str:
        """
        获取目录大小
        
        Args:
            path: 目录路径
        
        Returns:
            str: 格式化的大小字符串
        """
        total_size = 0
        
        for root, dirs, files in os.walk(path):
            for file in files:
                file_path = os.path.join(root, file)
                total_size += os.path.getsize(file_path)
        
        # 格式化大小
        if total_size < 1024:
            return f"{total_size}B"
        elif total_size < 1024 * 1024:
            return f"{total_size / 1024:.1f}KB"
        elif total_size < 1024 * 1024 * 1024:
            return f"{total_size / (1024 * 1024):.1f}MB"
        else:
            return f"{total_size / (1024 * 1024 * 1024):.1f}GB"
    
    def get_dir_creation_time(self, path: str) -> str:
        """
        获取目录创建时间
        
        Args:
            path: 目录路径
        
        Returns:
            str: 格式化的时间字符串
        """
        try:
            # 获取目录创建时间
            stat = os.stat(path)
            
            # 尝试获取创建时间，不同系统可能有不同的属性
            if hasattr(stat, 'st_birthtime'):  # macOS
                creation_time = stat.st_birthtime
            else:  # Linux
                creation_time = stat.st_mtime  # 使用修改时间作为创建时间
            
            return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(creation_time))
        except Exception:
            return "未知"
    
    def delete_model(self, model_path: str) -> bool:
        """
        删除本地模型
        
        Args:
            model_path: 模型完整路径
        
        Returns:
            bool: 是否删除成功
        """
        if not model_path or not os.path.exists(model_path):
            print(f"[DEBUG] 模型路径不存在: {model_path}")
            return False
        
        try:
            print(f"[DEBUG] 开始删除模型目录: {model_path}")
            shutil.rmtree(model_path)
            print(f"[DEBUG] 模型目录删除成功: {model_path}")
            return True
        except Exception as e:
            print(f"[DEBUG] 删除模型失败: {str(e)}")
            return False


# 测试代码
if __name__ == "__main__":
    # 直接在这里创建ModelManager实例，避免循环导入
    class TestModelManager:
        """测试用的模型管理器"""
        
        def list_models(self):
            """列出模型"""
            return [
                {"model_id": "test-model-1", "size": "1.2GB", "created_at": "2024-01-01 10:00:00"},
                {"model_id": "test-model-2", "size": "800MB", "created_at": "2024-01-02 14:30:00"}
            ]
    
    # 使用测试类
    manager = TestModelManager()
    
    # 测试列出模型
    print("测试列出模型...")
    models = manager.list_models()
    for model in models:
        print(f"模型: {model['model_id']}, 大小: {model['size']}, 创建时间: {model['created_at']}")
    
    print("\n注意: 这是一个简化的测试模式，完整功能需要通过app.py运行")