""" This node is a comfyui node for the moondream visual LLM (https://huggingface.co/vikhyatk/moondream2) It is customized through the following project: https://github.com/Hangover3832/ComfyUI-Hangover-Moondream """ from transformers import AutoModelForCausalLM as AutoModel # Try to import multiple possible tokenizers to improve compatibility try: from transformers import CodeGenTokenizerFast as Tokenizer except ImportError: try: from transformers import AutoTokenizer as Tokenizer print("[Moondream] Use AutoTokenizer instead of CodeGenTokenizerFast") except ImportError: from transformers import PreTrainedTokenizerFast as Tokenizer print("[Moondream] Use PreTrainedTokenizerFast instead of CodeGenTokenizerFast") from PIL import Image import torch import gc import numpy as np import codecs import subprocess import os import requests import sys import traceback def Run_git_status(repo:str) -> list[str]: """resturns a list of all model tag references for this huggingface repo""" url = f"https://huggingface.co/{repo}" process = subprocess.Popen(['git', 'ls-remote', url], stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr = process.communicate() result = [] if process.returncode == 0: revs = stdout.decode().splitlines() revs = [r.replace('main', 'latest') for r in revs if ('/tags' in r) or ('/main' in r)] for line in revs: rev = line.split('\t') result.append(f"{rev[-1].split('/')[-1]} -> {rev[0]}") return result class Moondream: HUGGINGFACE_MODEL_NAME = "vikhyatk/moondream2" DEVICES = ["gpu", "cpu"] if torch.cuda.is_available() else ["cpu"] Versions = 'versions.txt' Model_Revisions_URL = f"https://huggingface.co/{HUGGINGFACE_MODEL_NAME}/raw/main/{Versions}" current_path = os.path.abspath(os.path.dirname(__file__)) try: print("[Moondream] trying to update model versions...", end='') response = requests.get(Model_Revisions_URL) if response.status_code == 200: with open(f"{current_path}/{Versions}", 'w') as f: f.write(response.text) print('ok') except Exception as e: if hasattr(e, 'message'): msg = e.message else: msg = e print(f'failed ({msg})') with open(f"{current_path}/{Versions}", 'r') as f: versions = f.read() MODEL_REVISIONS = [v for v in versions.splitlines() if v.strip()] print(f"[Moondream] found model versions: {', '.join(MODEL_REVISIONS)}") MODEL_REVISIONS.insert(0,'ComfyUI/models/moondream2') try: print('\033[92m\033[4m[Moondream] model revsion references:\033[0m\033[92m') git_status = Run_git_status(HUGGINGFACE_MODEL_NAME) for s in git_status: print(s) # return ("",) except: pass finally: print('\033[0m') def __init__(self): self.model = None self.tokenizer = None self.revision = None @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "prompt": ("STRING", {"multiline": True, "default": "What is the gender, the age range, hair color, and hairstyle of the person in the image? Respond with only these four attributes."},), "separator": ("STRING", {"multiline": False, "default": r", "},), "model_revision": (s.MODEL_REVISIONS, {"default": s.MODEL_REVISIONS[-1]},), "temperature": ("FLOAT", {"min": 0.0, "max": 1.0, "step": 0.01, "default": 0.},), "device": (s.DEVICES, {"default": s.DEVICES[0]},) } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("text",) FUNCTION = "interrogate" OUTPUT_NODE = False CATEGORY = "InstantStudio" def interrogate(self, image:torch.Tensor, prompt:str, separator:str, model_revision:str, temperature:float, device:str): dev = "cuda" if device.lower() == "gpu" else "cpu" if temperature < 0.01: temperature = None do_sample = None else: do_sample = True if (self.model == None) or (self.tokenizer == None) or (device != self.device) or (model_revision != self.revision): del self.model del self.tokenizer gc.collect() if (device == "cpu") and torch.cuda.is_available(): torch.cuda.empty_cache() self.model = None self.tokenizer = None self.revision = model_revision print(f"[Moondream] loading model moondream2 revision '{model_revision}', please stand by....") if model_revision == Moondream.MODEL_REVISIONS[0]: model_name = model_revision model_revision = None else: model_name = Moondream.HUGGINGFACE_MODEL_NAME try: self.model = AutoModel.from_pretrained( model_name, trust_remote_code=True, revision=model_revision ).to(dev) try: # Try to load the tokenizer self.tokenizer = Tokenizer.from_pretrained(model_name) except Exception as e: # Try to load the tokenizer using AutoTokenizer if the default one fails from transformers import AutoTokenizer print(f"[Moondream] Failed to load tokenizer, trying to use AutoTokenizer... ({str(e)})") self.tokenizer = AutoTokenizer.from_pretrained(model_name) except Exception as e: import transformers error_msg = f"[Moondream] Failed to load model. transformers version: {transformers.__version__}\nError details: {str(e)}" print(error_msg) traceback.print_exc() raise ValueError(f"[Moondream] Please check if the transformers package meets the requirements. " "Note: older models may no longer be compatible with newer versions of the package.\n" f"Error details: {str(e)}") print(f"[Moondream] model moondream2 revision '{model_revision}' loaded successfully") self.device = device descriptions = "" prompts = list(filter(lambda x: x!="", [s.lstrip() for s in prompt.splitlines()])) # make a prompt list and remove unnecessary whitechars and empty lines if len(prompts) == 0: prompts = [""] try: for im in image: i = 255. * im.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) enc_image = self.model.encode_image(img) descr = "" sep = codecs.decode(separator, 'unicode_escape') for p in prompts: answer = self.model.answer_question(enc_image, p, self.tokenizer, temperature=temperature, do_sample=do_sample) descr += f"{answer}{sep}" descriptions += f"{descr[0:-len(sep)]}\n" except Exception as e: error_msg = f"[Moondream] Model execution failed: {str(e)}" print(error_msg) traceback.print_exc() raise ValueError(f"[Moondream] Model execution failed: {str(e)}") return(descriptions[0:-1],)