364 lines
14 KiB
Python
364 lines
14 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
import folder_paths
|
|
from PIL import Image
|
|
from huggingface_hub import hf_hub_download, snapshot_download
|
|
from transformers import AutoModelForCausalLM
|
|
|
|
# Register model_path
|
|
models_dir = os.path.join(folder_paths.base_path, "models")
|
|
ovis_dir = os.path.join(models_dir, "ovis")
|
|
os.makedirs(ovis_dir, exist_ok=True)
|
|
|
|
# Add Ovis models folder to folder_paths
|
|
if "ovis" not in folder_paths.folder_names_and_paths:
|
|
folder_paths.folder_names_and_paths["ovis"] = ([ovis_dir], folder_paths.supported_pt_extensions)
|
|
|
|
class Ovis2ModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": (["AIDC-AI/Ovis2-34B", "AIDC-AI/Ovis2-16B", "AIDC-AI/Ovis2-8B",
|
|
"AIDC-AI/Ovis2-2B", "AIDC-AI/Ovis2-1B"], {"default": "AIDC-AI/Ovis2-8B"}),
|
|
"precision": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
|
|
"max_token_length": ("INT", {"default": 32768, "min": 2048, "max": 65536}),
|
|
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
|
"auto_download": (["enable", "disable"], {"default": "enable"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("OVIS2_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "Ovis2"
|
|
|
|
def download_model(self, model_name):
|
|
"""Download the model files from Hugging Face if they don't exist locally."""
|
|
local_dir = os.path.join(ovis_dir, model_name.split('/')[-1])
|
|
|
|
print(f"Downloading Ovis2 model: {model_name} to {local_dir}")
|
|
try:
|
|
# Create a complete snapshot of the repository locally
|
|
snapshot_download(
|
|
repo_id=model_name,
|
|
local_dir=local_dir,
|
|
local_dir_use_symlinks=False # Use actual files instead of symlinks for better compatibility
|
|
)
|
|
print(f"Successfully downloaded {model_name} to {local_dir}")
|
|
return local_dir
|
|
except Exception as e:
|
|
print(f"Error downloading model: {str(e)}")
|
|
raise RuntimeError(f"Failed to download model {model_name}. Error: {str(e)}")
|
|
|
|
def check_model_files(self, model_name):
|
|
"""Check if the model files already exist locally."""
|
|
local_dir = os.path.join(ovis_dir, model_name.split('/')[-1])
|
|
|
|
# Check for config.json as a basic indicator that the model exists
|
|
config_path = os.path.join(local_dir, "config.json")
|
|
return os.path.exists(config_path), local_dir
|
|
|
|
def load_model(self, model_name, precision, max_token_length, device, auto_download):
|
|
print(f"Loading Ovis2 model: {model_name}")
|
|
|
|
# Set precision
|
|
if precision == "bfloat16":
|
|
dtype = torch.bfloat16
|
|
elif precision == "float16":
|
|
dtype = torch.float16
|
|
else:
|
|
dtype = torch.float32
|
|
|
|
# Check if model exists locally
|
|
model_exists, local_dir = self.check_model_files(model_name)
|
|
|
|
# Download model if it doesn't exist and auto_download is enabled
|
|
if not model_exists and auto_download == "enable":
|
|
self.download_model(model_name)
|
|
|
|
# Load the model
|
|
try:
|
|
# First try loading from local directory
|
|
if model_exists or auto_download == "enable":
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
local_dir,
|
|
torch_dtype=dtype,
|
|
multimodal_max_length=max_token_length,
|
|
trust_remote_code=True
|
|
).to(device)
|
|
except Exception as e:
|
|
print(f"Error loading from local directory, falling back to HuggingFace: {str(e)}")
|
|
# Fall back to loading directly from HuggingFace
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_name,
|
|
torch_dtype=dtype,
|
|
multimodal_max_length=max_token_length,
|
|
trust_remote_code=True
|
|
).to(device)
|
|
else:
|
|
# Load directly from HuggingFace
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_name,
|
|
torch_dtype=dtype,
|
|
multimodal_max_length=max_token_length,
|
|
trust_remote_code=True
|
|
).to(device)
|
|
|
|
# Get tokenizers
|
|
text_tokenizer = model.get_text_tokenizer()
|
|
visual_tokenizer = model.get_visual_tokenizer()
|
|
|
|
return ({"model": model, "text_tokenizer": text_tokenizer, "visual_tokenizer": visual_tokenizer},)
|
|
except Exception as e:
|
|
print(f"Error loading Ovis2 model: {str(e)}")
|
|
raise e
|
|
|
|
|
|
class Ovis2ImageCaption:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("OVIS2_MODEL",),
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"default": "Describe this image in detail.", "multiline": True}),
|
|
"max_new_tokens": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"temperature": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("caption",)
|
|
FUNCTION = "generate_caption"
|
|
CATEGORY = "Ovis2"
|
|
|
|
def generate_caption(self, model, image, prompt, max_new_tokens, temperature):
|
|
model_data = model
|
|
model = model_data["model"]
|
|
text_tokenizer = model_data["text_tokenizer"]
|
|
|
|
# Convert ComfyUI image to PIL image format
|
|
# ComfyUI images are in format [batch, height, width, channel]
|
|
if len(image.shape) == 4:
|
|
# Handle batch of images - take the first one
|
|
image_tensor = image[0]
|
|
else:
|
|
# Handle single image
|
|
image_tensor = image
|
|
|
|
pil_image = Image.fromarray((image_tensor * 255).astype(np.uint8))
|
|
|
|
# Process the image with Ovis2
|
|
query = f"<image>\n{prompt}"
|
|
|
|
try:
|
|
# Preprocess inputs
|
|
_, input_ids, pixel_values = model.preprocess_inputs(query, [pil_image], max_partition=9)
|
|
|
|
# Generate response
|
|
outputs = model.generate(
|
|
input_ids=input_ids.to(model.device),
|
|
pixel_values=pixel_values.to(model.device),
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
top_p=0.9,
|
|
do_sample=True
|
|
)
|
|
|
|
# Decode the generated text
|
|
generated_text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
# Extract just the generated response (remove the instruction)
|
|
response = generated_text[len(query):]
|
|
|
|
return (response,)
|
|
except Exception as e:
|
|
print(f"Error generating caption: {str(e)}")
|
|
return (f"Error generating caption: {str(e)}",)
|
|
|
|
|
|
class Ovis2MultiImageInput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("OVIS2_MODEL",),
|
|
"image1": ("IMAGE",),
|
|
"prompt": ("STRING", {"default": "Analyze these images. What are the similarities and differences?", "multiline": True}),
|
|
"max_new_tokens": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"temperature": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.1}),
|
|
},
|
|
"optional": {
|
|
"image2": ("IMAGE",),
|
|
"image3": ("IMAGE",),
|
|
"image4": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("response",)
|
|
FUNCTION = "process_multi_images"
|
|
CATEGORY = "Ovis2"
|
|
|
|
def process_multi_images(self, model, image1, prompt, max_new_tokens, temperature, image2=None, image3=None, image4=None):
|
|
model_data = model
|
|
model = model_data["model"]
|
|
text_tokenizer = model_data["text_tokenizer"]
|
|
|
|
# Convert ComfyUI images to PIL image format
|
|
images = []
|
|
|
|
# Process first image
|
|
if len(image1.shape) == 4:
|
|
# Handle batch of images - take the first one
|
|
pil_image1 = Image.fromarray((image1[0] * 255).astype(np.uint8))
|
|
else:
|
|
# Handle single image
|
|
pil_image1 = Image.fromarray((image1 * 255).astype(np.uint8))
|
|
|
|
images.append(pil_image1)
|
|
|
|
# Add optional additional images if provided
|
|
if image2 is not None:
|
|
if len(image2.shape) == 4:
|
|
pil_image2 = Image.fromarray((image2[0] * 255).astype(np.uint8))
|
|
else:
|
|
pil_image2 = Image.fromarray((image2 * 255).astype(np.uint8))
|
|
images.append(pil_image2)
|
|
|
|
if image3 is not None:
|
|
if len(image3.shape) == 4:
|
|
pil_image3 = Image.fromarray((image3[0] * 255).astype(np.uint8))
|
|
else:
|
|
pil_image3 = Image.fromarray((image3 * 255).astype(np.uint8))
|
|
images.append(pil_image3)
|
|
|
|
if image4 is not None:
|
|
if len(image4.shape) == 4:
|
|
pil_image4 = Image.fromarray((image4[0] * 255).astype(np.uint8))
|
|
else:
|
|
pil_image4 = Image.fromarray((image4 * 255).astype(np.uint8))
|
|
images.append(pil_image4)
|
|
|
|
# Process the images with Ovis2
|
|
query = f"<image>\n{prompt}"
|
|
|
|
try:
|
|
# Preprocess inputs
|
|
_, input_ids, pixel_values = model.preprocess_inputs(query, images, max_partition=9)
|
|
|
|
# Generate response
|
|
outputs = model.generate(
|
|
input_ids=input_ids.to(model.device),
|
|
pixel_values=pixel_values.to(model.device),
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
top_p=0.9,
|
|
do_sample=True
|
|
)
|
|
|
|
# Decode the generated text
|
|
generated_text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
# Extract just the generated response (remove the instruction)
|
|
response = generated_text[len(query):]
|
|
|
|
return (response,)
|
|
except Exception as e:
|
|
print(f"Error processing images: {str(e)}")
|
|
return (f"Error processing images: {str(e)}",)
|
|
|
|
|
|
class Ovis2VideoFramesDescription:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("OVIS2_MODEL",),
|
|
"frames": ("IMAGE",),
|
|
"prompt": ("STRING", {"default": "Describe what's happening in this video.", "multiline": True}),
|
|
"max_new_tokens": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"temperature": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.1}),
|
|
"frame_skip": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}),
|
|
"max_frames": ("INT", {"default": 16, "min": 1, "max": 32, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("description",)
|
|
FUNCTION = "describe_video"
|
|
CATEGORY = "Ovis2"
|
|
|
|
def describe_video(self, model, frames, prompt, max_new_tokens, temperature, frame_skip, max_frames):
|
|
model_data = model
|
|
model = model_data["model"]
|
|
text_tokenizer = model_data["text_tokenizer"]
|
|
|
|
# Convert ComfyUI image frames to PIL images
|
|
pil_images = []
|
|
|
|
# Check if frames is a single image or a batch
|
|
if len(frames.shape) == 4 and frames.shape[0] > 1:
|
|
# It's a batch of images
|
|
# Apply frame skipping and limit to max_frames
|
|
selected_indices = list(range(0, min(frames.shape[0], max_frames * frame_skip), frame_skip))
|
|
|
|
for i in selected_indices:
|
|
if i < frames.shape[0]:
|
|
pil_image = Image.fromarray((frames[i] * 255).astype(np.uint8))
|
|
pil_images.append(pil_image)
|
|
else:
|
|
# It's a single image, treat it as a 1-frame video
|
|
if len(frames.shape) == 4:
|
|
pil_image = Image.fromarray((frames[0] * 255).astype(np.uint8))
|
|
else:
|
|
pil_image = Image.fromarray((frames * 255).astype(np.uint8))
|
|
pil_images.append(pil_image)
|
|
|
|
# Process the images with Ovis2
|
|
query = f"<image>\n{prompt}"
|
|
|
|
try:
|
|
# Preprocess inputs
|
|
_, input_ids, pixel_values = model.preprocess_inputs(query, pil_images, max_partition=9)
|
|
|
|
# Generate response
|
|
outputs = model.generate(
|
|
input_ids=input_ids.to(model.device),
|
|
pixel_values=pixel_values.to(model.device),
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
top_p=0.9,
|
|
do_sample=True
|
|
)
|
|
|
|
# Decode the generated text
|
|
generated_text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
# Extract just the generated response (remove the instruction)
|
|
response = generated_text[len(query):]
|
|
|
|
return (response,)
|
|
except Exception as e:
|
|
print(f"Error describing video: {str(e)}")
|
|
return (f"Error describing video: {str(e)}",)
|
|
|
|
|
|
# Register the nodes
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Ovis2ModelLoader": Ovis2ModelLoader,
|
|
"Ovis2ImageCaption": Ovis2ImageCaption,
|
|
"Ovis2MultiImageInput": Ovis2MultiImageInput,
|
|
"Ovis2VideoFramesDescription": Ovis2VideoFramesDescription,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Ovis2ModelLoader": "Load Ovis2 Model",
|
|
"Ovis2ImageCaption": "Ovis2 Image Caption",
|
|
"Ovis2MultiImageInput": "Ovis2 Multi-Image Analysis",
|
|
"Ovis2VideoFramesDescription": "Ovis2 Video Frames Description",
|
|
}
|