filliptm-comfyui-fl-penguinvl/fl_utils/model_manager.py

"""
Model loading and caching for FL PenguinVL.
Handles downloading, caching, and loading of Penguin-VL models.
"""

import gc
import os
import sys
import importlib.util
from pathlib import Path
from typing import Dict, Any, Optional, Callable

import torch

# Import paths module explicitly to avoid relative import issues
_FL_UTILS_DIR = Path(__file__).parent


def _import_paths():
    module_path = _FL_UTILS_DIR / "paths.py"
    spec = importlib.util.spec_from_file_location("penguinvl_paths", str(module_path))
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


_paths = _import_paths()
get_models_directory = _paths.get_models_directory

# Global model cache
_MODEL_CACHE: Dict[str, Any] = {}

# Model variant configurations
MODEL_VARIANTS = {
    "2B": {
        "hf_repo": "tencent/Penguin-VL-2B",
        "description": "2B parameter model - Fast, lower VRAM",
        "vram_fp16": 6,
        "vram_4bit": 3,
    },
    "8B": {
        "hf_repo": "tencent/Penguin-VL-8B",
        "description": "8B parameter model - Higher quality, more VRAM",
        "vram_fp16": 18,
        "vram_4bit": 8,
    },
}


def get_variant_list() -> list:
    """Get list of available model variants."""
    return list(MODEL_VARIANTS.keys())


def get_variant_info(variant: str) -> dict:
    """Get info for a specific variant."""
    return MODEL_VARIANTS.get(variant, MODEL_VARIANTS["2B"])


def download_model_if_needed(variant: str, force: bool = False) -> Path:
    """
    Download model from HuggingFace if not present.

    Args:
        variant: Model variant (2B or 8B)
        force: Force re-download

    Returns:
        Path to model directory
    """
    from huggingface_hub import snapshot_download

    models_dir = get_models_directory()
    variant_dir = models_dir / f"Penguin-VL-{variant}"

    if variant_dir.exists() and not force:
        if (variant_dir / "config.json").exists():
            print(f"[FL PenguinVL] Model Penguin-VL-{variant} already downloaded")
            return variant_dir

    hf_repo = MODEL_VARIANTS[variant]["hf_repo"]
    print(f"[FL PenguinVL] Downloading {hf_repo}...")
    snapshot_download(
        repo_id=hf_repo,
        local_dir=str(variant_dir),
        local_dir_use_symlinks=False,
        ignore_patterns=["*.md", ".git*"],
    )
    print(f"[FL PenguinVL] Download complete: Penguin-VL-{variant}")
    return variant_dir


def load_model(
    variant: str = "2B",
    precision: str = "auto",
    device_str: str = "auto",
    use_4bit: bool = False,
    force_reload: bool = False,
    progress_callback: Optional[Callable[[int, int], None]] = None,
) -> dict:
    """
    Load Penguin-VL model + processor with caching.

    Args:
        variant: Model variant (2B or 8B)
        precision: Precision mode (auto, fp32, fp16, bf16)
        device_str: Device selection (auto, cuda, cpu)
        use_4bit: Whether to use 4-bit quantization
        force_reload: Force reload even if cached
        progress_callback: Optional callback(current_step, total_steps) for progress updates

    Returns:
        Dict with: model, processor, device, dtype, variant
    """
    # Determine device
    if device_str == "auto":
        if torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")
    else:
        device = torch.device(device_str)

    # Determine dtype
    # bf16 is required for the vision encoder (fp16 overflows in flash attention)
    if precision == "auto":
        if device.type == "cuda":
            if torch.cuda.is_bf16_supported():
                dtype = torch.bfloat16
            else:
                dtype = torch.float16
        else:
            dtype = torch.float32
    elif precision == "fp16":
        dtype = torch.float16
    elif precision == "bf16":
        dtype = torch.bfloat16
    else:
        dtype = torch.float32

    # Create cache key
    cache_key = f"{variant}_{device}_{dtype}_{use_4bit}"

    # Return cached model if available
    if not force_reload and cache_key in _MODEL_CACHE:
        print(f"[FL PenguinVL] Using cached model: Penguin-VL-{variant}")
        return _MODEL_CACHE[cache_key]

    # Clear existing cache before loading (single-model policy for large VLMs)
    if _MODEL_CACHE:
        print("[FL PenguinVL] Clearing previous model from cache...")
        clear_model_cache()

    # Stage 1: Download if needed
    model_path = download_model_if_needed(variant, force=force_reload)
    if progress_callback:
        progress_callback(1, 3)

    from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor

    # Build model kwargs
    model_kwargs = {
        "dtype": dtype,
        "trust_remote_code": True,
    }

    # Setup quantization if requested
    if use_4bit:
        if device.type != "cuda":
            print(f"[FL PenguinVL] WARNING: 4-bit quantization requires CUDA, but device is {device.type}")
            print("[FL PenguinVL] Disabling 4-bit quantization, using full precision instead")
            use_4bit = False
        else:
            try:
                from transformers import BitsAndBytesConfig
                model_kwargs["quantization_config"] = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_compute_dtype=dtype,
                    bnb_4bit_use_double_quant=True,
                    bnb_4bit_quant_type="nf4",
                )
                print("[FL PenguinVL] Using 4-bit quantization")
            except ImportError:
                print("[FL PenguinVL] WARNING: bitsandbytes not installed, using full precision")
                use_4bit = False

    # Stage 2: Load model
    # Penguin-VL's __init__ calls from_pretrained for its vision encoder, but
    # transformers wraps __init__ in torch.device("meta") context which crashes
    # nested from_pretrained calls. Fix: pre-load config and set vision_encoder=None
    # to force the _from_config path. The vision encoder weights (311 keys) are in
    # the main checkpoint and get loaded by the outer from_pretrained regardless.
    print(f"[FL PenguinVL] Loading Penguin-VL-{variant} on {device} ({dtype})...")
    config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True)
    config.vision_encoder = None

    model = AutoModelForCausalLM.from_pretrained(str(model_path), config=config, **model_kwargs)
    if not use_4bit:
        model = model.to(device)
    model.eval()
    if progress_callback:
        progress_callback(2, 3)

    # Stage 3: Load processor
    processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True)
    if progress_callback:
        progress_callback(3, 3)

    model_info = {
        "model": model,
        "processor": processor,
        "device": device,
        "dtype": dtype,
        "variant": variant,
    }

    _MODEL_CACHE[cache_key] = model_info
    print(f"[FL PenguinVL] Penguin-VL-{variant} loaded successfully!")
    return model_info


def clear_model_cache():
    """Clear the model cache and free GPU memory."""
    global _MODEL_CACHE

    for key in list(_MODEL_CACHE.keys()):
        model_info = _MODEL_CACHE.pop(key)
        if "model" in model_info:
            del model_info["model"]
        if "processor" in model_info:
            del model_info["processor"]

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    print("[FL PenguinVL] Model cache cleared")