nodes/1038lab-ComfyUI-MiniCPM
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2025-08-24 00:45:56 -07:00
committed by GitHub
parent e9b87230b4
commit 0768ca9d6d
2 changed files with 462 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AILab_MiniCPM.py
+388
View File
@@ -0,0 +1,388 @@
import os
import torch
import folder_paths
from transformers import AutoTokenizer, AutoModel
from torchvision.transforms.v2 import ToPILImage
from comfy.comfy_types import IO
from comfy_api.input import VideoInput
import json
import gc
import sys
import io
from pathlib import Path
os.environ['TRANSFORMERS_VERBOSITY'] = 'error'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
if hasattr(torch.backends, 'cuda'):
if hasattr(torch.backends.cuda, 'matmul'):
torch.backends.cuda.matmul.allow_tf32 = True
if hasattr(torch.backends.cuda, 'allow_tf32'):
torch.backends.cuda.allow_tf32 = True
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
# Load configuration file
with open(Path(__file__).parent / "minicpm_config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODEL_SETTINGS = config["model_settings"]
PROMPT_TYPES = config.get("prompt_types", {})
TRANSFORMERS_MODELS = config.get("transformers_models", {})
_MODEL_CACHE = {}
class MiniCPM_Transformers_Models:
def __init__(self, model: str, processing_mode: str):
try:
models_dir = Path(folder_paths.models_dir).resolve()
prompt_generator_dir = (models_dir / "LLM").resolve()
prompt_generator_dir.mkdir(parents=True, exist_ok=True)
model_config = TRANSFORMERS_MODELS[model]
model_id = model_config["name"]
self.model_checkpoint = prompt_generator_dir / Path(model_id).name
if not self.model_checkpoint.exists():
print(f"Downloading model: {model_id}")
from huggingface_hub import snapshot_download
snapshot_download(
repo_id=model_id,
local_dir=str(self.model_checkpoint),
local_dir_use_symlinks=False,
)
self.device = torch.device("cuda" if processing_mode == "GPU" and torch.cuda.is_available() else "cpu")
self.bf16_support = (
torch.cuda.is_available()
and torch.cuda.get_device_capability(self.device)[0] >= 8
)
old_stdout = sys.stdout
old_stderr = sys.stderr
try:
sys.stdout = io.StringIO()
sys.stderr = io.StringIO()
self.tokenizer = AutoTokenizer.from_pretrained(
str(self.model_checkpoint),
trust_remote_code=True,
low_cpu_mem_usage=True,
)
self.model = AutoModel.from_pretrained(
str(self.model_checkpoint),
trust_remote_code=True,
low_cpu_mem_usage=True,
attn_implementation="sdpa",
torch_dtype=torch.bfloat16 if self.bf16_support else torch.float16,
)
if processing_mode == "GPU" and torch.cuda.is_available():
self.model = self.model.to(self.device)
self.model.eval()
finally:
sys.stdout = old_stdout
sys.stderr = old_stderr
except Exception as e:
raise RuntimeError(f"Model initialization failed: {str(e)}") from e
def generate(self, images, system: str, prompt: str, max_new_tokens: int,
temperature: float, top_p: float, top_k: int,
repetition_penalty: float, video_max_slice_nums: int = 2, seed: int = -1) -> str:
try:
if seed != -1:
torch.manual_seed(seed)
with torch.no_grad():
if isinstance(images, list):
# Video frame processing
msgs = [{"role": "user", "content": images + [prompt]}]
else:
# Single image processing
msgs = [{"role": "user", "content": [images, prompt]}]
params = {"use_image_id": False, "max_slice_nums": video_max_slice_nums}
result = self.model.chat(
image=None,
msgs=msgs,
tokenizer=self.tokenizer,
sampling=True,
top_k=top_k,
top_p=top_p,
temperature=temperature,
repetition_penalty=repetition_penalty,
max_new_tokens=max_new_tokens,
**params,
)
return result.strip()
except Exception as e:
return f"Generation error: {str(e)}"
finally:
gc.collect()
class MiniCPM_Transformers_Base:
def __init__(self):
self.predictor = None
self.current_processing_mode = None
self.current_model = None
def _load_model(self, model: str, processing_mode: str, memory_management: str):
cache_key = f"{model}_{processing_mode}"
try:
if memory_management == "Global Cache":
if cache_key in _MODEL_CACHE:
self.predictor = _MODEL_CACHE[cache_key]
else:
self.predictor = MiniCPM_Transformers_Models(model, processing_mode)
_MODEL_CACHE[cache_key] = self.predictor
elif (self.predictor is None or self.current_processing_mode != processing_mode or self.current_model != model):
if self.predictor is not None:
del self.predictor
self.predictor = None
torch.cuda.empty_cache()
gc.collect()
self.predictor = MiniCPM_Transformers_Models(model, processing_mode)
self.current_processing_mode = processing_mode
self.current_model = model
except Exception as model_error:
raise RuntimeError(f"Model loading error: {str(model_error)}")
def _cleanup_memory(self, memory_management: str):
if memory_management == "Clear After Run":
try:
del self.predictor
self.predictor = None
torch.cuda.empty_cache()
gc.collect()
except:
pass
def _process_image(self, image):
try:
if isinstance(image, (list, tuple)) and len(image) > 0:
img_tensor = image[0]
else:
img_tensor = image
if not isinstance(img_tensor, torch.Tensor):
raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}")
if img_tensor.dim() == 4:
img_tensor = img_tensor[0]
elif img_tensor.dim() == 2:
img_tensor = img_tensor.unsqueeze(0).repeat(3, 1, 1)
if img_tensor.shape[0] == 3:
pass
elif img_tensor.shape[-1] == 3:
img_tensor = img_tensor.permute(2, 0, 1)
else:
raise ValueError(f"Unexpected image tensor shape: {img_tensor.shape}")
pil_image = ToPILImage()(img_tensor).convert("RGB")
return pil_image
except Exception as input_error:
raise RuntimeError(f"Input processing error: {str(input_error)}")
def _generate_response(self, images, system_prompt, prompt, **gen_params):
try:
with torch.inference_mode():
response = self.predictor.generate(
images=images,
system=system_prompt,
prompt=prompt,
**gen_params
)
return response
except Exception as gen_error:
raise RuntimeError(f"Generation error: {str(gen_error)}")
def encode_video(self, source_video, MAX_NUM_FRAMES):
def uniform_sample(l, n):
gap = len(l) / n
idxs = [int(i * gap + gap / 2) for i in range(n)]
return [l[i] for i in idxs]
components = source_video.get_components()
vr = components.images
avg_fps = float(components.frame_rate)
sample_fps = round(avg_fps / 1)
frame_idx = [i for i in range(0, len(vr), sample_fps)]
if len(frame_idx) > MAX_NUM_FRAMES:
frame_idx = uniform_sample(frame_idx, MAX_NUM_FRAMES)
frames = [vr[idx] for idx in frame_idx]
frames = [ToPILImage()(v.permute([2, 0, 1])).convert("RGB") for v in frames]
return frames
class AILab_MiniCPM_4_V(MiniCPM_Transformers_Base):
@classmethod
def INPUT_TYPES(cls):
model_list = list(TRANSFORMERS_MODELS.keys())
return {
"required": {},
"optional": {
"image": ("IMAGE",),
"video": ("VIDEO",),
"model": (model_list, {"default": model_list[0]}),
"preset_prompt": (list(PROMPT_TYPES.keys()), {"default": "Describe"}),
"custom_prompt": ("STRING", {"default": "", "multiline": True}),
"device": (["Auto", "GPU", "CPU"], {"default": "Auto"}),
"memory_management": (["Keep in Memory", "Clear After Run", "Global Cache"], {"default": "Keep in Memory"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("STRING",)
FUNCTION = "generate"
CATEGORY = "🧪AILab/MiniCPM"
def generate(self, image=None, video=None, model=None, preset_prompt="Describe", custom_prompt="", device="Auto", memory_management="Keep in Memory", seed=-1):
try:
if image is None and video is None:
return ("Error: Please provide either an image or video input.",)
pm = ("GPU" if torch.cuda.is_available() else "CPU") if device == "Auto" else device
model = model or list(TRANSFORMERS_MODELS.keys())[0]
self._load_model(model, pm, memory_management)
if video is not None:
frames = self.encode_video(video, 64)
if frames:
images = frames
else:
return ("Error: No frames extracted from video.",)
elif image is not None:
pil_image = self._process_image(image)
images = pil_image
else:
return ("Error: No valid input provided.",)
preset = PROMPT_TYPES.get(preset_prompt, "Describe this image.")
prompt = custom_prompt.strip() if (isinstance(custom_prompt, str) and custom_prompt.strip()) else preset
if isinstance(seed, int) and seed != -1:
torch.manual_seed(seed)
response = self._generate_response(
images=images,
system_prompt=MODEL_SETTINGS.get("default_system_prompt", ""),
prompt=prompt,
max_new_tokens=MODEL_SETTINGS["default_max_tokens"],
temperature=MODEL_SETTINGS["default_temperature"],
top_p=MODEL_SETTINGS["default_top_p"],
top_k=MODEL_SETTINGS["default_top_k"],
repetition_penalty=MODEL_SETTINGS["default_repetition_penalty"],
video_max_slice_nums=2,
seed=seed,
)
self._cleanup_memory(memory_management)
return (response,)
except Exception as e:
self._cleanup_memory(memory_management)
return (f"Error: {str(e)}",)
class AILab_MiniCPM_4_V_Advanced(MiniCPM_Transformers_Base):
@classmethod
def INPUT_TYPES(cls):
model_list = list(TRANSFORMERS_MODELS.keys())
return {
"required": {},
"optional": {
"image": ("IMAGE",),
"video": ("VIDEO",),
"model": (model_list, {"default": model_list[0]}),
"preset_prompt": (list(PROMPT_TYPES.keys()), {"default": "Describe"}),
"custom_prompt": ("STRING", {"default": "", "multiline": True}),
"system_prompt": ("STRING", {"default": MODEL_SETTINGS.get("default_system_prompt", ""), "multiline": True}),
"max_new_tokens": ("INT", {"default": MODEL_SETTINGS["default_max_tokens"], "min": 1, "max": 4096}),
"temperature": ("FLOAT", {"default": MODEL_SETTINGS["default_temperature"], "min": 0, "max": 2.0, "step": 0.05}),
"top_p": ("FLOAT", {"default": MODEL_SETTINGS["default_top_p"], "min": 0.0, "max": 1.0, "step": 0.01}),
"top_k": ("INT", {"default": MODEL_SETTINGS["default_top_k"], "min": 0, "max": 200}),
"repetition_penalty": ("FLOAT", {"default": MODEL_SETTINGS["default_repetition_penalty"], "min": 0.8, "max": 1.5, "step": 0.01}),
"video_max_num_frames": ("INT", {"default": 64, "min": 1, "max": 128}),
"video_max_slice_nums": ("INT", {"default": 2, "min": 1, "max": 4}),
"device": (["Auto", "GPU", "CPU"], {"default": "Auto"}),
"memory_management": (["Keep in Memory", "Clear After Run", "Global Cache"], {"default": "Keep in Memory"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
}
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("PROMPT", "STRING")
FUNCTION = "generate"
CATEGORY = "🧪AILab/MiniCPM"
def generate(self, image=None, video=None, model=None, preset_prompt="Describe", custom_prompt="", system_prompt="", max_new_tokens=None, temperature=None, top_p=None, top_k=None, repetition_penalty=None, video_max_num_frames=64, video_max_slice_nums=2, device="Auto", memory_management="Keep in Memory", seed=-1):
try:
if image is None and video is None:
return ("", "Error: Please provide either an image or video input.")
pm = ("GPU" if torch.cuda.is_available() else "CPU") if device == "Auto" else device
model = model or list(TRANSFORMERS_MODELS.keys())[0]
self._load_model(model, pm, memory_management)
if video is not None:
frames = self.encode_video(video, video_max_num_frames)
if frames:
images = frames
else:
return ("", "Error: No frames extracted from video.")
elif image is not None:
pil_image = self._process_image(image)
images = pil_image
else:
return ("", "Error: No valid input provided.")
preset = PROMPT_TYPES.get(preset_prompt, "Describe this image.")
prompt = custom_prompt.strip() if (isinstance(custom_prompt, str) and custom_prompt.strip()) else preset
if isinstance(seed, int) and seed != -1:
torch.manual_seed(seed)
max_new_tokens = max_new_tokens if max_new_tokens is not None else MODEL_SETTINGS["default_max_tokens"]
temperature = temperature if temperature is not None else MODEL_SETTINGS["default_temperature"]
top_p = top_p if top_p is not None else MODEL_SETTINGS["default_top_p"]
top_k = top_k if top_k is not None else MODEL_SETTINGS["default_top_k"]
repetition_penalty = repetition_penalty if repetition_penalty is not None else MODEL_SETTINGS["default_repetition_penalty"]
response = self._generate_response(
images=images,
system_prompt=system_prompt or MODEL_SETTINGS.get("default_system_prompt", ""),
prompt=prompt,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
top_k=top_k,
repetition_penalty=repetition_penalty,
video_max_slice_nums=video_max_slice_nums,
seed=seed,
)
self._cleanup_memory(memory_management)
return (prompt, response)
except Exception as e:
self._cleanup_memory(memory_management)
return ("", f"Error: {str(e)}")
NODE_CLASS_MAPPINGS = {
"AILab_MiniCPM_4_V": AILab_MiniCPM_4_V,
"AILab_MiniCPM_4_V_Advanced": AILab_MiniCPM_4_V_Advanced,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AILab_MiniCPM_4_V": "MiniCPM-4-V",
"AILab_MiniCPM_4_V_Advanced": "MiniCPM-4-V Advanced",
}
__init__.py
+74
View File
@@ -0,0 +1,74 @@
import os
import importlib
import sys
from pathlib import Path
current_dir = Path(__file__).parent
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
EXCLUDE_FILES = ['__init__.py', '__pycache__']
# Enable Windows color support
if sys.platform == 'win32':
os.system('color')
# Check if llama-cpp-python is available for GGUF functionality
GGUF_AVAILABLE = False
try:
import llama_cpp
GGUF_AVAILABLE = True
except ImportError:
# Use Windows color codes for better visibility
print("\n" + "=" * 80)
print("\033[91mWARNING: llama-cpp-python library not found, GGUF functionality is not available\033[0m")
print("\033[93mTo use GGUF features, install additional dependencies:\033[0m")
print("\033[96mpip install llama-cpp-python\033[0m")
# Check if installation guide exists and provide link
install_guide = current_dir / "llama_cpp_install.md"
if install_guide.exists():
print("\033[93mFor detailed installation instructions with CUDA support, please see:\033[0m")
print(f"\033[96m{install_guide}\033[0m")
print("\033[92mBasic MiniCPM functionality is still available\033[0m")
print("=" * 80 + "\n")
except Exception as e:
print("\n" + "=" * 80)
print(f"\033[91mError loading GGUF dependencies: {str(e)}\033[0m")
print("\033[92mBasic MiniCPM functionality is still available\033[0m")
print("=" * 80 + "\n")
# Process all Python files in the directory (auto-registration functionality)
for file in current_dir.glob('*.py'):
if file.name not in EXCLUDE_FILES:
try:
module_name = file.stem
spec = importlib.util.spec_from_file_location(module_name, str(file))
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
# Skip GGUF module if llama-cpp-python is not available
if not GGUF_AVAILABLE and 'GGUF' in module_name:
print(f"\033[93mSkipping {module_name} - GGUF functionality not available\033[0m")
continue
spec.loader.exec_module(module)
if hasattr(module, 'NODE_CLASS_MAPPINGS'):
NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
print(f"\033[92mLoaded {module_name} nodes: {list(module.NODE_CLASS_MAPPINGS.keys())}\033[0m")
if hasattr(module, 'NODE_DISPLAY_NAME_MAPPINGS'):
NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
except Exception as e:
print(f"\033[91mError loading module {module_name}: {str(e)}\033[0m")
if 'GGUF' not in module_name: # Only show warning for non-GGUF modules
print(f"\033[93mWarning: Failed to load {module_name} module\033[0m")
print(f"\n\033[92mMiniCPM nodes loaded: {list(NODE_CLASS_MAPPINGS.keys())}\033[0m")
print(f"\033[92mTotal nodes registered: {len(NODE_CLASS_MAPPINGS)}\033[0m")
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']