wTechArtist-ComfyUI_VVL_VideoCamera_Advanced/vggt_model_loader.py

"""
VGGT Model Loader
独立的VGGT模型加载和管理模块
"""

import os
import logging
from typing import Dict, Optional
import torch

# ComfyUI相关导入
try:
    import folder_paths
    import comfy.model_management
    COMFYUI_AVAILABLE = True
except ImportError:
    folder_paths = None
    comfy = None
    COMFYUI_AVAILABLE = False

# VGGT相关导入
try:
    from vggt.models.vggt import VGGT
    VGGT_AVAILABLE = True
except Exception as e:
    VGGT = None
    VGGT_AVAILABLE = False
    _VGGT_IMPORT_ERROR = e

# 配置日志
logger = logging.getLogger('vvl_vggt_loader')

# VGGT模型配置
VGGT_MODEL_DIR_NAME = "vggt"
VGGT_MODEL_CONFIG = {
    "VGGT-1B": {
        "model_name": "facebook/VGGT-1B", 
        "description": "VGGT 1B parameter model for camera pose estimation",
        "file_name": "vggt_1b.pt",
        "size_mb": 4700  # 约4.7GB
    },
    # 可以添加更多模型配置
    # "VGGT-Large": {
    #     "model_name": "facebook/VGGT-Large",
    #     "description": "VGGT Large model with better accuracy",
    #     "file_name": "vggt_large.pt"
    # }
}

# 全局模型缓存
_VGGT_MODEL_CACHE = {}

def get_vggt_model_dir() -> str:
    """获取VGGT模型目录路径"""
    if COMFYUI_AVAILABLE and folder_paths:
        # 使用ComfyUI的models目录
        model_dir = os.path.join(folder_paths.models_dir, VGGT_MODEL_DIR_NAME)
    else:
        # 备用路径
        model_dir = os.path.join(os.path.expanduser("~"), ".cache", "comfyui", "models", VGGT_MODEL_DIR_NAME)
    
    # 确保目录存在
    os.makedirs(model_dir, exist_ok=True)
    return model_dir

def get_vggt_model_path(model_config: dict) -> Optional[str]:
    """获取VGGT模型的本地路径"""
    model_dir = get_vggt_model_dir()
    model_path = os.path.join(model_dir, model_config["file_name"])
    
    # 检查模型文件是否存在
    if os.path.exists(model_path):
        logger.info(f"Found local VGGT model: {model_path}")
        return model_path
    
    logger.warning(f"Local VGGT model not found: {model_path}")
    logger.info(f"Expected model directory: {model_dir}")
    return None

def list_available_models() -> Dict[str, bool]:
    """列出可用的模型及其本地可用性"""
    available_models = {}
    for model_name, config in VGGT_MODEL_CONFIG.items():
        local_path = get_vggt_model_path(config)
        available_models[model_name] = local_path is not None
    return available_models

def load_vggt_model(model_name: str, device: torch.device) -> Optional[torch.nn.Module]:
    """加载VGGT模型"""
    if not VGGT_AVAILABLE:
        logger.error(f"VGGT not available: {_VGGT_IMPORT_ERROR}")
        return None
    
    try:
        # 检查缓存
        cache_key = f"{model_name}_{device}"
        if cache_key in _VGGT_MODEL_CACHE:
            logger.info(f"Using cached VGGT model: {model_name}")
            return _VGGT_MODEL_CACHE[cache_key]
        
        if model_name not in VGGT_MODEL_CONFIG:
            raise ValueError(f"Unknown VGGT model: {model_name}. Available: {list(VGGT_MODEL_CONFIG.keys())}")
        
        config = VGGT_MODEL_CONFIG[model_name]
        local_path = get_vggt_model_path(config)
        
        # 加载模型
        if local_path and os.path.exists(local_path):
            logger.info(f"Loading VGGT model from local path: {local_path}")
            # 注意：这里可能需要根据实际的本地加载API调整
            # 目前先使用HuggingFace的方式，后续可以优化为直接加载本地权重
            model = VGGT.from_pretrained(config["model_name"])
        else:
            # 从HuggingFace加载（会自动下载）
            logger.info(f"Loading VGGT model from HuggingFace: {config['model_name']}")
            logger.warning(f"Model will be downloaded (~{config.get('size_mb', 'Unknown')}MB)")
            model = VGGT.from_pretrained(config["model_name"])
        
        # 移动到指定设备并设置为评估模式
        model = model.to(device)
        model.eval()
        
        # 缓存模型
        _VGGT_MODEL_CACHE[cache_key] = model
        
        logger.info(f"VGGT model {model_name} loaded successfully on {device}")
        return model
        
    except Exception as e:
        logger.error(f"Failed to load VGGT model {model_name}: {e}")
        return None

def clear_model_cache():
    """清除模型缓存"""
    global _VGGT_MODEL_CACHE
    for key in list(_VGGT_MODEL_CACHE.keys()):
        del _VGGT_MODEL_CACHE[key]
    logger.info("VGGT model cache cleared")

def get_model_info() -> Dict:
    """获取模型信息"""
    return {
        "available": VGGT_AVAILABLE,
        "error": _VGGT_IMPORT_ERROR if not VGGT_AVAILABLE else None,
        "model_dir": get_vggt_model_dir(),
        "cached_models": list(_VGGT_MODEL_CACHE.keys()),
        "config": VGGT_MODEL_CONFIG,
        "local_models": list_available_models()
    }

class VVLVGGTLoader:
    """VGGT模型加载器节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        vggt_models = list(VGGT_MODEL_CONFIG.keys())
        device_list = ["auto", "cuda", "cpu"]
        
        return {
            "required": {
                "device": (device_list, {
                    "default": "auto",
                    "tooltip": "选择运行设备。auto会自动选择CUDA（如果可用）或CPU"
                }),
                "vggt_model": (vggt_models, {
                    "default": vggt_models[0] if vggt_models else "VGGT-1B",
                    "tooltip": "选择VGGT模型版本。VGGT-1B是标准的10亿参数模型，约4.7GB"
                }),
            }
        }

    RETURN_TYPES = ("VVL_VGGT_MODEL",)
    RETURN_NAMES = ("vggt_model",)
    FUNCTION = "load_vggt_model"
    CATEGORY = "💃VVL/VGGT"

    def load_vggt_model(self, device: str, vggt_model: str):
        """加载VGGT模型并返回模型实例"""
        
        # 确定设备
        if device == "auto":
            torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            torch_device = torch.device(device)
        
        logger.info(f"VVLVGGTLoader: Loading VGGT model: {vggt_model} on device: {torch_device}")
        
        # 检查VGGT是否可用
        if not VGGT_AVAILABLE:
            error_msg = f"VGGT not available: {_VGGT_IMPORT_ERROR}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)
        
        # 显示模型信息
        model_info = get_model_info()
        logger.info(f"Model directory: {model_info['model_dir']}")
        logger.info(f"Local models: {model_info['local_models']}")
        
        # 加载模型
        model_instance = load_vggt_model(vggt_model, torch_device)
        
        if model_instance is None:
            raise RuntimeError(f"Failed to load VGGT model: {vggt_model}")
        
        # 创建包含模型和相关信息的字典
        model_data = {
            'model': model_instance,
            'device': torch_device,
            'model_name': vggt_model,
            'config': VGGT_MODEL_CONFIG[vggt_model],
            'loader_info': model_info
        }
        
        logger.info("VVLVGGTLoader: VGGT model loaded successfully")
        return (model_data,)