411 lines
17 KiB
Python
411 lines
17 KiB
Python
from transformers import (
|
|
AutoTokenizer,
|
|
AutoModelForSequenceClassification,
|
|
AutoModelForCausalLM,
|
|
AutoModelForSeq2SeqLM,
|
|
AutoConfig,
|
|
BitsAndBytesConfig
|
|
)
|
|
from llama_cpp import Llama
|
|
import torch
|
|
import os
|
|
import folder_paths
|
|
import re
|
|
import subprocess
|
|
import datetime
|
|
import shutil
|
|
|
|
GLOBAL_MODELS_DIR = os.path.join(folder_paths.models_dir, "LLM_checkpoints")
|
|
|
|
WEB_DIRECTORY = "./web/assets/js"
|
|
|
|
class AnyType(str):
|
|
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
|
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
|
|
class LLM_Node:
|
|
def __init__(self, device="cuda"):
|
|
self.device = device
|
|
self.custom_nodes_folder = folder_paths.folder_names_and_paths['custom_nodes'][0][0]
|
|
self.comfy_ui_llm_node_path = os.path.join(self.custom_nodes_folder, "ComfyUI_LLM_Node")
|
|
# Check if bfloat16 is supported by the device
|
|
self.supports_bfloat16 = 'cuda' in device and torch.cuda.is_bf16_supported()
|
|
self.fixmeused = False
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Get a list of directories in the checkpoints_path
|
|
model_options = []
|
|
for name in os.listdir(GLOBAL_MODELS_DIR):
|
|
dir_path = os.path.join(GLOBAL_MODELS_DIR, name)
|
|
if os.path.isdir(dir_path):
|
|
if "GGUF" in name:
|
|
gguf_files = [os.path.join(name, file) for file in os.listdir(dir_path) if file.endswith('.gguf')]
|
|
model_options.extend(gguf_files)
|
|
else:
|
|
model_options.append(name)
|
|
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True, "default": ""}),
|
|
"seed": ("INT", {"default": 777}),
|
|
"model": (model_options, ),
|
|
"max_tokens": ("INT", {"default": 2000, "min": 1}),
|
|
"apply_chat_template": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"AdvOptionsConfig": ("ADVOPTIONSCONFIG",),
|
|
"QuantizationConfig": ("QUANTIZATIONCONFIG",),
|
|
"CodingConfig": ("CODINGCONFIG",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("string",)
|
|
OUTPUT_NODE = False
|
|
FUNCTION = "main"
|
|
CATEGORY = "LLM"
|
|
|
|
def fixme(self, filename, tokenizer, model_to_use, generate_kwargs, apply_chat_template, text):
|
|
self.fixmeused = True
|
|
generated_text = "\n\n**************************************** FIX ME loop ****************************************\n\n"
|
|
generated_text += text+"\n\n\n"
|
|
|
|
file_path = os.path.join(self.comfy_ui_llm_node_path, "tmp", filename)
|
|
|
|
try:
|
|
generated_text += self.generate_text(text, tokenizer, model_to_use, generate_kwargs, apply_chat_template)
|
|
except Exception as e:
|
|
print(f"Failed to generate new text for {filename}: {e}")
|
|
return
|
|
|
|
match = re.search(r'```python\s*([\s\S]+?)\s*```', generated_text)
|
|
if not match:
|
|
print(f"No code block found in generated text for {filename}.")
|
|
return
|
|
|
|
new_code = match.group(1)
|
|
|
|
# Write the new content back to the file
|
|
try:
|
|
with open(file_path, 'w') as file:
|
|
file.write(new_code)
|
|
print(f"File {filename} has been updated successfully.")
|
|
return generated_text
|
|
except IOError as e:
|
|
print(f"Failed to write new content to file {filename}: {e}")
|
|
return
|
|
|
|
def extract_files_and_code(self, text):
|
|
files_code = []
|
|
# First pattern to check for filenames and associated Python code
|
|
primary_pattern = r'\*\*(.+?\.py):\*\*\s*```python\s*([\s\S]+?)```'
|
|
matches = re.findall(primary_pattern, text)
|
|
|
|
if matches:
|
|
for filename, code in matches:
|
|
files_code.append((filename, code.strip()))
|
|
else:
|
|
# Secondary pattern to check for Python code blocks without filenames
|
|
secondary_pattern = r'```python\s*([\s\S]+?)\s*```'
|
|
code_blocks = re.findall(secondary_pattern, text)
|
|
for code in code_blocks:
|
|
files_code.append(("main.py", code.strip()))
|
|
|
|
return files_code
|
|
|
|
def write_files_to_folder(self, files_code, folder_path):
|
|
for filename, code in files_code:
|
|
file_path = os.path.join(folder_path, filename)
|
|
with open(file_path, 'w') as file:
|
|
file.write(code)
|
|
|
|
def log_history(self, user_text, generated_text):
|
|
history_dir = os.path.join(self.comfy_ui_llm_node_path, "history")
|
|
timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
filename = datetime.datetime.now().strftime('%Y-%m-%d') + '.txt'
|
|
filepath = os.path.join(history_dir, filename)
|
|
with open(filepath, 'a') as file:
|
|
file.write(f"{timestamp} - User: {user_text}\n")
|
|
file.write(f"{timestamp} - Generated: {generated_text}\n\n")
|
|
|
|
def generate_text(self, text, tokenizer, model_to_use, generate_kwargs, apply_chat_template):
|
|
if apply_chat_template:
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": text}
|
|
]
|
|
text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True
|
|
)
|
|
|
|
input_ids = tokenizer([text], return_tensors="pt").input_ids.to(self.device)
|
|
outputs = model_to_use.generate(input_ids, **generate_kwargs)
|
|
|
|
if apply_chat_template:
|
|
generated_ids = [output_ids[len(input_id):] for input_id, output_ids in zip(input_ids, outputs)]
|
|
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
|
else:
|
|
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
return generated_text
|
|
|
|
def main(self, text, seed, model, max_tokens, apply_chat_template, AdvOptionsConfig=None, QuantizationConfig=None, CodingConfig=None):
|
|
model_path = os.path.join(GLOBAL_MODELS_DIR, model)
|
|
generated_text = None
|
|
if "GGUF" in model:
|
|
generate_kwargs = {'max_tokens': max_tokens}
|
|
|
|
if AdvOptionsConfig:
|
|
for option in ['temperature', 'top_p', 'top_k', 'repetition_penalty']:
|
|
if option in AdvOptionsConfig:
|
|
if (option == 'repetition_penalty'):
|
|
option1 = 'repeat_penalty'
|
|
else:
|
|
option1 = option
|
|
generate_kwargs[option1] = AdvOptionsConfig[option]
|
|
|
|
|
|
model_to_use = Llama(
|
|
model_path=model_path,
|
|
n_gpu_layers=-1,
|
|
seed=seed,
|
|
# n_ctx=2048, # Uncomment to increase the context window
|
|
)
|
|
generated_text = model_to_use(text, **generate_kwargs)
|
|
self.log_history(text, generated_text['choices'][0]['text'])
|
|
return (generated_text['choices'][0]['text'],)
|
|
else:
|
|
torch.manual_seed(seed)
|
|
|
|
model_kwargs = {
|
|
'device_map': 'auto',
|
|
'quantization_config': QuantizationConfig
|
|
}
|
|
|
|
if AdvOptionsConfig:
|
|
if 'trust_remote_code' in AdvOptionsConfig:
|
|
model_kwargs['trust_remote_code'] = AdvOptionsConfig['trust_remote_code']
|
|
|
|
if 'torch_dtype' in AdvOptionsConfig and hasattr(torch, AdvOptionsConfig['torch_dtype']):
|
|
model_kwargs['torch_dtype'] = getattr(torch, AdvOptionsConfig['torch_dtype'])
|
|
|
|
config = AutoConfig.from_pretrained(model_path, **model_kwargs)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
|
|
if config.model_type == "t5":
|
|
model_to_use = AutoModelForSeq2SeqLM.from_pretrained(model_path, **model_kwargs)
|
|
elif config.model_type in ["gpt2", "gpt_refact", "gemma", "llama", "mistral", "qwen2"]:
|
|
model_to_use = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
|
|
elif config.model_type == "bert":
|
|
model_to_use = AutoModelForSequenceClassification.from_pretrained(model_path, **model_kwargs)
|
|
else:
|
|
raise ValueError(f"Unsupported model type: {config.model_type}")
|
|
|
|
generate_kwargs = {'max_length': max_tokens}
|
|
|
|
if AdvOptionsConfig:
|
|
for option in ['temperature', 'top_p', 'top_k', 'repetition_penalty']:
|
|
if option in AdvOptionsConfig:
|
|
generate_kwargs[option] = AdvOptionsConfig[option]
|
|
|
|
if config.model_type in ["t5", "gpt2", "gpt_refact", "gemma", "llama", "mistral", "qwen2"]:
|
|
generated_text = self.generate_text(text, tokenizer, model_to_use, generate_kwargs, apply_chat_template)
|
|
chat_story = generated_text
|
|
if CodingConfig and CodingConfig.get('execute_code'):
|
|
files_code = self.extract_files_and_code(generated_text)
|
|
execution_attempts = 0
|
|
while execution_attempts < 10:
|
|
if files_code:
|
|
tmp_folder_path = os.path.join(self.comfy_ui_llm_node_path, "tmp")
|
|
if self.fixmeused == False:
|
|
os.makedirs(tmp_folder_path, exist_ok=True)
|
|
self.write_files_to_folder(files_code, tmp_folder_path)
|
|
pycache_path = os.path.join(tmp_folder_path, "__pycache__")
|
|
if os.path.exists(pycache_path):
|
|
shutil.rmtree(pycache_path)
|
|
files = os.listdir(tmp_folder_path)
|
|
files.sort(key=lambda x: x == 'main.py')
|
|
for filename in files:
|
|
file_path = os.path.join(tmp_folder_path, filename)
|
|
try:
|
|
env = os.environ.copy()
|
|
env["SDL_VIDEODRIVER"] = "dummy"
|
|
command = ['python', file_path]
|
|
try:
|
|
subprocess.run(command, capture_output=True, text=True, check=True, env=env, timeout=1) # no output only check for errors
|
|
print(f"Successful Execution: {filename}")
|
|
except subprocess.TimeoutExpired:
|
|
print("Execution timed out after 1 second. Terminating subprocess.")
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Execution failed, retrying... Error: {e.stderr}")
|
|
with open(file_path, 'r') as file:
|
|
file_data = file.read()
|
|
text = f"Error encountered: {e.stderr}\nFix the code. Write all code. Don't take shortcut. Don't write 'same as your original code'!\n{file_data}"
|
|
chat_story += self.fixme(filename, tokenizer, model_to_use, generate_kwargs, apply_chat_template, text)
|
|
execution_attempts += 1
|
|
break
|
|
else: # no more files to check
|
|
break
|
|
else: # no code found
|
|
break
|
|
self.log_history(text, chat_story)
|
|
if (CodingConfig.get('project_folder')):
|
|
for filename in files:
|
|
source_file_path = os.path.join(tmp_folder_path, filename)
|
|
destination_file_path = os.path.join(CodingConfig.get('project_folder'), filename)
|
|
shutil.copy(source_file_path, destination_file_path)
|
|
main_py_path = os.path.join(CodingConfig.get('project_folder'), 'main.py')
|
|
subprocess.run(['python', main_py_path], capture_output=True, text=True, check=True)
|
|
if os.path.exists(tmp_folder_path):
|
|
shutil.rmtree(tmp_folder_path)
|
|
return (chat_story,)
|
|
else:
|
|
self.log_history(text, chat_story)
|
|
return (chat_story,)
|
|
elif config.model_type == "bert":
|
|
return ("BERT model detected; specific task handling not implemented in this example.",)
|
|
|
|
class Output_Node:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": (any, {}),
|
|
}
|
|
}
|
|
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "main"
|
|
CATEGORY = "LLM"
|
|
RETURN_TYPES = ()
|
|
|
|
def main(self, text):
|
|
return {"ui": {"text": (text,)}}
|
|
|
|
class QuantizationConfig_Node:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
quantization_modes = ["none", "load_in_8bit", "load_in_4bit"]
|
|
return {
|
|
"required": {
|
|
"quantization_mode": (quantization_modes, {"default": "none"}),
|
|
"llm_int8_threshold": ("FLOAT", {"default": 6.0}),
|
|
"llm_int8_skip_modules": ("STRING", {"default": ""}),
|
|
"llm_int8_enable_fp32_cpu_offload": ("BOOLEAN", {"default": False}),
|
|
"llm_int8_has_fp16_weight": ("BOOLEAN", {"default": False}),
|
|
"bnb_4bit_compute_dtype": ("STRING", {"default": "float32"}),
|
|
"bnb_4bit_quant_type": ("STRING", {"default": "fp4"}),
|
|
"bnb_4bit_use_double_quant": ("BOOLEAN", {"default": False}),
|
|
"bnb_4bit_quant_storage": ("STRING", {"default": "uint8"}),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "main"
|
|
CATEGORY = "LLM"
|
|
RETURN_TYPES = ("QUANTIZATIONCONFIG",)
|
|
RETURN_NAMES = ("QuantizationConfig",)
|
|
|
|
def main(self, quantization_mode, llm_int8_threshold: float = 6.0, llm_int8_skip_modules="", llm_int8_enable_fp32_cpu_offload=False, llm_int8_has_fp16_weight=False, bnb_4bit_compute_dtype="float32", bnb_4bit_quant_type="fp4", bnb_4bit_use_double_quant=False, bnb_4bit_quant_storage="uint8"):
|
|
|
|
llm_int8_skip_modules_list = llm_int8_skip_modules.split(',') if llm_int8_skip_modules else []
|
|
|
|
quantization_config = BitsAndBytesConfig(
|
|
load_in_8bit=quantization_mode == "load_in_8bit",
|
|
load_in_4bit=quantization_mode == "load_in_4bit",
|
|
llm_int8_threshold=float(llm_int8_threshold),
|
|
llm_int8_skip_modules=llm_int8_skip_modules_list,
|
|
llm_int8_enable_fp32_cpu_offload=llm_int8_enable_fp32_cpu_offload,
|
|
llm_int8_has_fp16_weight=llm_int8_has_fp16_weight,
|
|
bnb_4bit_compute_dtype=getattr(torch, bnb_4bit_compute_dtype, torch.float32),
|
|
bnb_4bit_quant_type=bnb_4bit_quant_type,
|
|
bnb_4bit_use_double_quant=bnb_4bit_use_double_quant,
|
|
bnb_4bit_quant_storage=bnb_4bit_quant_storage,
|
|
)
|
|
|
|
return (quantization_config,)
|
|
|
|
class AdvOptionsNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
dtype_options = ["auto", "float32", "bfloat16", "float16", "float64"]
|
|
return {
|
|
"required": {
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.9, "min": 0.1, "step": 0.1}),
|
|
"top_k": ("INT", {"default": 50, "min": 0}),
|
|
"repetition_penalty": ("FLOAT", {"default": 1.2, "min": 0.1, "step": 0.1}),
|
|
"trust_remote_code": ("BOOLEAN", {"default": False}),
|
|
"torch_dtype": (dtype_options, {"default": "auto"}),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "main"
|
|
CATEGORY = "LLM"
|
|
RETURN_TYPES = ("ADVOPTIONSCONFIG",)
|
|
RETURN_NAMES = ("AdvOptionsConfig",)
|
|
|
|
def main(self, temperature=1.0, top_p=0.9, top_k=50, repetition_penalty=1.2, trust_remote_code=False, torch_dtype="auto"):
|
|
options_config = {
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"repetition_penalty": repetition_penalty,
|
|
"trust_remote_code": trust_remote_code,
|
|
"torch_dtype": torch_dtype,
|
|
}
|
|
|
|
return (options_config,)
|
|
|
|
class CodingOptionsNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"execute_code": ("BOOLEAN", {"default": False}),
|
|
"project_folder": ("STRING", {"default": "/projects/LLM"}),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "main"
|
|
CATEGORY = "LLM"
|
|
RETURN_TYPES = ("CODINGCONFIG",)
|
|
RETURN_NAMES = ("CodingConfig",)
|
|
|
|
def main(self, execute_code, project_folder):
|
|
|
|
return ({"execute_code":execute_code,"project_folder":project_folder},)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LLM_Node": LLM_Node,
|
|
"Output_Node": Output_Node,
|
|
"QuantizationConfig_Node": QuantizationConfig_Node,
|
|
"AdvOptions_Node": AdvOptionsNode,
|
|
"CodingOptionsNode": CodingOptionsNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LLM_Node": "LLM Node",
|
|
"Output_Node": "Output Node",
|
|
"QuantizationConfig_Node": "Quantization Config Node",
|
|
"AdvOptions_Node": "Advanced Options Node",
|
|
"CodingOptionsNode": "Code Config Node",
|
|
}
|