FillipandClaude Opus 4.7
3c04bdc186
Adds two nodes (FL VLM Model Loader, FL VLM Image Inference) with a Qwen3-VL backend, optional 8-bit/4-bit quantization, resize presets, per-image or shared-prompt batch modes, and a streaming progress UI. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
112 lines
3.2 KiB
Python
112 lines
3.2 KiB
Python
import gc
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import AutoProcessor, BitsAndBytesConfig, Qwen3VLForConditionalGeneration
|
|
|
|
import folder_paths
|
|
|
|
from .model_registry import dtype_from_name
|
|
|
|
|
|
logger = logging.getLogger("FL_VLM")
|
|
|
|
_MODEL_CACHE = {}
|
|
|
|
|
|
def get_vlm_models_dir():
|
|
model_dir = Path(folder_paths.models_dir) / "VLM"
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
return model_dir
|
|
|
|
|
|
def register_vlm_folder():
|
|
model_dir = str(get_vlm_models_dir())
|
|
supported = {".safetensors", ".bin", ".json", ".txt", ".model"}
|
|
if "VLM" not in folder_paths.folder_names_and_paths:
|
|
folder_paths.folder_names_and_paths["VLM"] = ([model_dir], supported)
|
|
return
|
|
|
|
paths, extensions = folder_paths.folder_names_and_paths["VLM"]
|
|
if model_dir not in paths:
|
|
paths.append(model_dir)
|
|
extensions.update(supported)
|
|
|
|
|
|
def _safe_model_dir_name(model_id):
|
|
return model_id.replace("/", "--")
|
|
|
|
|
|
def _quantization_config(quantization):
|
|
if quantization == "4bit":
|
|
return BitsAndBytesConfig(load_in_4bit=True)
|
|
if quantization == "8bit":
|
|
return BitsAndBytesConfig(load_in_8bit=True)
|
|
return None
|
|
|
|
|
|
def _device_map(device_map):
|
|
if device_map == "cuda":
|
|
return {"": "cuda:0"}
|
|
if device_map == "cpu":
|
|
return {"": "cpu"}
|
|
return device_map
|
|
|
|
|
|
def clear_model_cache():
|
|
_MODEL_CACHE.clear()
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def load_vlm(model_id, dtype, quantization, attention, device_map, force_reload=False):
|
|
register_vlm_folder()
|
|
cache_key = (model_id, dtype, quantization, attention, device_map)
|
|
|
|
if force_reload:
|
|
_MODEL_CACHE.pop(cache_key, None)
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
if cache_key in _MODEL_CACHE:
|
|
logger.info("Using cached VLM model: %s", model_id)
|
|
return _MODEL_CACHE[cache_key]
|
|
|
|
model_path = get_vlm_models_dir() / _safe_model_dir_name(model_id)
|
|
if not model_path.exists() or not any(model_path.iterdir()):
|
|
logger.info("Downloading VLM model %s to %s", model_id, model_path)
|
|
snapshot_download(repo_id=model_id, local_dir=str(model_path), local_dir_use_symlinks=False)
|
|
|
|
torch_dtype = dtype_from_name(dtype)
|
|
quant_config = _quantization_config(quantization)
|
|
|
|
load_kwargs = {
|
|
"torch_dtype": torch_dtype,
|
|
"device_map": _device_map(device_map),
|
|
"attn_implementation": attention,
|
|
}
|
|
if quant_config is not None:
|
|
load_kwargs["quantization_config"] = quant_config
|
|
|
|
logger.info("Loading VLM model from %s", model_path)
|
|
model = Qwen3VLForConditionalGeneration.from_pretrained(str(model_path), **load_kwargs)
|
|
processor = AutoProcessor.from_pretrained(str(model_path))
|
|
|
|
model_bundle = {
|
|
"backend": "qwen3_vl",
|
|
"model": model,
|
|
"processor": processor,
|
|
"model_id": model_id,
|
|
"model_path": str(model_path),
|
|
"dtype": dtype,
|
|
"quantization": quantization,
|
|
"attention": attention,
|
|
"device_map": device_map,
|
|
}
|
|
_MODEL_CACHE[cache_key] = model_bundle
|
|
return model_bundle
|