├── .gitattributes ├── .gitignore ├── __init__.py ├── install.py ├── nodes.py └── requirements.txt /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | checkpoints/ 3 | *.py[cod] 4 | *$py.class 5 | *.egg-info 6 | .pytest_cache -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS 2 | 3 | __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] -------------------------------------------------------------------------------- /install.py: -------------------------------------------------------------------------------- 1 | # from https://github.com/gokayfem/ComfyUI_VLM_nodes/blob/main/install_init.py 2 | 3 | import platform 4 | import subprocess 5 | import sys 6 | import importlib.util 7 | import re 8 | import torch 9 | import packaging.tags 10 | from requests import get 11 | 12 | def get_python_version(): 13 | """Return the Python version in a concise format, e.g., '39' for Python 3.9.""" 14 | version_match = re.match(r"3\.(\d+)", platform.python_version()) 15 | if version_match: 16 | return "3" + version_match.group(1) 17 | else: 18 | return None 19 | 20 | def get_system_info(): 21 | """Gather system information related to NVIDIA GPU, CUDA version, AVX2 support, Python version, OS, and platform tag.""" 22 | system_info = { 23 | 'gpu': False, 24 | 'cuda_version': None, 25 | 'avx2': False, 26 | 'python_version': get_python_version(), 27 | 'os': platform.system(), 28 | 'os_bit': platform.architecture()[0].replace("bit", ""), 29 | 'platform_tag': None, 30 | } 31 | 32 | # Check for NVIDIA GPU and CUDA version 33 | if importlib.util.find_spec('torch'): 34 | system_info['gpu'] = torch.cuda.is_available() 35 | if system_info['gpu']: 36 | system_info['cuda_version'] = "cu" + torch.version.cuda.replace(".", "").strip() 37 | 38 | # Check for AVX2 support 39 | if importlib.util.find_spec('cpuinfo'): 40 | try: 41 | # Attempt to import the cpuinfo module 42 | import cpuinfo 43 | 44 | # Safely attempt to retrieve CPU flags 45 | cpu_info = cpuinfo.get_cpu_info() 46 | if cpu_info and 'flags' in cpu_info: 47 | # Check if 'avx2' is among the CPU flags 48 | system_info['avx2'] = 'avx2' in cpu_info['flags'] 49 | else: 50 | # Handle the case where CPU info is unavailable or does not contain 'flags' 51 | system_info['avx2'] = False 52 | except Exception as e: 53 | # Handle unexpected errors gracefully 54 | print(f"Error retrieving CPU information: {e}") 55 | system_info['avx2'] = False 56 | else: 57 | # Handle the case where the cpuinfo module is not installed 58 | print("cpuinfo module not available.") 59 | system_info['avx2'] = False 60 | # Determine the platform tag 61 | if importlib.util.find_spec('packaging.tags'): 62 | system_info['platform_tag'] = next(packaging.tags.sys_tags()).platform 63 | 64 | return system_info 65 | 66 | def latest_lamacpp(): 67 | try: 68 | response = get("https://api.github.com/repos/abetlen/llama-cpp-python/releases/latest") 69 | return response.json()["tag_name"].replace("v", "") 70 | except Exception: 71 | return "0.2.26" 72 | 73 | def install_package(package_name, custom_command=None): 74 | if not package_is_installed(package_name): 75 | print(f"Installing {package_name}...") 76 | command = [sys.executable, "-m", "pip", "install", package_name, "--no-cache-dir"] 77 | if custom_command: 78 | command += custom_command.split() 79 | subprocess.check_call(command) 80 | else: 81 | print(f"{package_name} is already installed.") 82 | 83 | def package_is_installed(package_name): 84 | return importlib.util.find_spec(package_name) is not None 85 | 86 | def install_llama(system_info): 87 | imported = package_is_installed("llama-cpp-python") or package_is_installed("llama_cpp") 88 | if imported: 89 | print("llama-cpp installed") 90 | else: 91 | lcpp_version = latest_lamacpp() 92 | base_url = "https://github.com/abetlen/llama-cpp-python/releases/download/v" 93 | avx = "AVX2" if system_info['avx2'] else "AVX" 94 | if system_info['gpu']: 95 | cuda_version = system_info['cuda_version'] 96 | custom_command = f"--force-reinstall --no-deps --index-url=https://jllllll.github.io/llama-cpp-python-cuBLAS-wheels/{avx}/{cuda_version}" 97 | else: 98 | custom_command = f"{base_url}{lcpp_version}/llama_cpp_python-{lcpp_version}-{system_info['platform_tag']}.whl" 99 | install_package("llama-cpp-python", custom_command=custom_command) 100 | 101 | def main(): 102 | system_info = get_system_info() 103 | install_llama(system_info) 104 | 105 | if __name__ == "__main__": 106 | main() -------------------------------------------------------------------------------- /nodes.py: -------------------------------------------------------------------------------- 1 | import os 2 | from llama_cpp import Llama 3 | 4 | import comfy.model_management as mm 5 | import comfy.utils 6 | import folder_paths 7 | llm_extensions = ['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.gguf'] 8 | 9 | script_directory = os.path.dirname(os.path.abspath(__file__)) 10 | 11 | folder_paths.folder_names_and_paths["LLM"] = ([os.path.join(folder_paths.models_dir, "LLM")], llm_extensions) 12 | 13 | class llama_cpp_model_loader: 14 | 15 | @classmethod 16 | def INPUT_TYPES(s): 17 | return {"required": { 18 | "n_gpu_layers": ("INT", {"default": 0, "min": -1, "max": 4096, "step": 1}), 19 | "download_default": ("BOOLEAN", {"default": False}), 20 | }, 21 | "optional": { 22 | "model": (folder_paths.get_filename_list("LLM"),), 23 | } 24 | } 25 | 26 | RETURN_TYPES = ("LLAMACPPMODEL",) 27 | RETURN_NAMES = ("llamamodel",) 28 | FUNCTION = "loadmodel" 29 | CATEGORY = "Llama-cpp" 30 | 31 | def loadmodel(self, n_gpu_layers, download_default, model=None): 32 | mm.soft_empty_cache() 33 | 34 | custom_config = { 35 | "model": model, 36 | } 37 | if not hasattr(self, "model") or custom_config != self.current_config: 38 | self.current_config = custom_config 39 | llama3_dir = (os.path.join(folder_paths.models_dir, 'LLM', 'llama3')) 40 | default_checkpoint_path = os.path.join(llama3_dir, 'Meta-Llama-3-8B-Instruct-Q4_K_M.gguf') 41 | default_model = "Meta-Llama-3-8B-Instruct-Q4_K_M.gguf" 42 | if download_default and not os.path.exists(default_checkpoint_path): 43 | print(f"Downloading {default_model}") 44 | from huggingface_hub import snapshot_download 45 | allow_patterns = [f"*{default_model}*"] 46 | snapshot_download(repo_id="bartowski/Meta-Llama-3-8B-Instruct-GGUF", 47 | allow_patterns=allow_patterns, 48 | local_dir=llama3_dir, 49 | local_dir_use_symlinks=False 50 | ) 51 | model_path = default_checkpoint_path 52 | else: 53 | model_path = os.path.join(folder_paths.models_dir, 'LLM', model) 54 | print(f"Loading model from {model_path}") 55 | 56 | llm = Llama(model_path, n_gpu_layers=n_gpu_layers) 57 | 58 | return (llm,) 59 | 60 | class llama_cpp_instruct: 61 | @classmethod 62 | def INPUT_TYPES(s): 63 | return {"required": { 64 | "llmamamodel": ("LLAMACPPMODEL",), 65 | "parameters": ("LLAMACPPARAMS", ), 66 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 67 | "prompt": ("STRING", {"multiline": True, "default": "How much wood would woodchuck chuck, if woodchuck could chuck wood?",}), 68 | 69 | }, 70 | } 71 | 72 | RETURN_TYPES = ("STRING",) 73 | RETURN_NAMES = ("output",) 74 | FUNCTION = "process" 75 | CATEGORY = "Llama-cpp" 76 | 77 | def process(self, llmamamodel, prompt, parameters, seed): 78 | 79 | mm.soft_empty_cache() 80 | output = llmamamodel( 81 | f"Q: {prompt} A: ", # Prompt 82 | stop=["Q:", "\n"], # Stop generating just before the model would generate a new question 83 | echo=False, 84 | seed=seed, 85 | max_tokens = parameters.get("max_tokens", 32), 86 | top_k = parameters.get("top_k", 40), 87 | top_p = parameters.get("top_p", 0.95), 88 | min_p = parameters.get("min_p", 0.05), 89 | typical_p = parameters.get("typical_p", 1.0), 90 | temperature = parameters.get("temperature", 0.8), 91 | repeat_penalty = parameters.get("repeat_penalty", 1.1), 92 | frequency_penalty = parameters.get("frequency_penalty", 0.0), 93 | presence_penalty = parameters.get("presence_penalty", 0.0), 94 | tfs_z = parameters.get("tfs_z", 1.0), 95 | mirostat_mode = parameters.get("mirostat_mode", 0), 96 | mirostat_eta = parameters.get("mirostat_eta", 0.1), 97 | mirostat_tau = parameters.get("mirostat_tau", 5.0), 98 | ) 99 | print(output) 100 | text = output['choices'][0]['text'] 101 | return (text,) 102 | 103 | class llama_cpp_parameters: 104 | @classmethod 105 | def INPUT_TYPES(s): 106 | return {"required": { 107 | "max_tokens": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 1}), 108 | "top_k": ("INT", {"default": 40, "min": 0, "max": 1000, "step": 1}), 109 | "top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}), 110 | "min_p": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01}), 111 | "typical_p": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), 112 | "temperature": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}), 113 | "repeat_penalty": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.01}), 114 | "frequency_penalty": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), 115 | "presence_penalty": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), 116 | "tfs_z": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), 117 | "mirostat_mode": ("INT", {"default": 0, "min": 0, "max": 1, "step": 1}), 118 | "mirostat_eta": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}), 119 | "mirostat_tau": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 10.0, "step": 0.01}), 120 | } 121 | } 122 | 123 | RETURN_TYPES = ("LLAMACPPARAMS",) 124 | RETURN_NAMES = ("parameters",) 125 | FUNCTION = "process" 126 | CATEGORY = "Llama-cpp" 127 | 128 | def process(self, max_tokens, top_k, top_p, min_p, typical_p, temperature, repeat_penalty, 129 | frequency_penalty, presence_penalty, tfs_z, mirostat_mode, mirostat_eta, mirostat_tau, 130 | ): 131 | 132 | parameters_dict = { 133 | "max_tokens": max_tokens, 134 | "top_k": top_k, 135 | "top_p": top_p, 136 | "min_p": min_p, 137 | "typical_p": typical_p, 138 | "temperature": temperature, 139 | "repeat_penalty": repeat_penalty, 140 | "frequency_penalty": frequency_penalty, 141 | "presence_penalty": presence_penalty, 142 | "tfs_z": tfs_z, 143 | "mirostat_mode": mirostat_mode, 144 | "mirostat_eta": mirostat_eta, 145 | "mirostat_tau": mirostat_tau, 146 | } 147 | return (parameters_dict,) 148 | 149 | NODE_CLASS_MAPPINGS = { 150 | "llama_cpp_model_loader": llama_cpp_model_loader, 151 | "llama_cpp_instruct": llama_cpp_instruct, 152 | "llama_cpp_parameters": llama_cpp_parameters 153 | } 154 | 155 | NODE_DISPLAY_NAME_MAPPINGS = { 156 | "llama_cpp_model_loader": "Llama-cpp Model Loader", 157 | "llama_cpp_instruct": "Llama-cpp Instruct", 158 | "llama_cpp_parameters": "Llama-cpp Parameters" 159 | } 160 | 161 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | diskcache --------------------------------------------------------------------------------