snapshot - testing loading LLM+vision from pretrained
This commit is contained in:
Olivv
+54
-34
@@ -15,7 +15,7 @@ but without frame prediction.
|
||||
"""
|
||||
import os
|
||||
import torch
|
||||
from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
|
||||
from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM, AutoModelForCausalLM, AutoTokenizer, AutoConfig
|
||||
from safetensors.torch import load_file
|
||||
import comfy.model_management as mm
|
||||
import comfy.sd as sd
|
||||
@@ -73,7 +73,7 @@ def list_clip_files():
|
||||
model_dir = os.path.join(mm.model_path, "clip")
|
||||
return sorted(f for f in os.listdir(model_dir) if f.endswith(".safetensors"))
|
||||
|
||||
class FunPackDualCLIPLoaderInstruct:
|
||||
class FunPackCLIPLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
base = nodes.DualCLIPLoader.INPUT_TYPES()
|
||||
@@ -81,11 +81,14 @@ class FunPackDualCLIPLoaderInstruct:
|
||||
'required':
|
||||
{
|
||||
'clip_model_name': (s.get_filename_list(),),
|
||||
'llama_instruct_model_name': (s.get_filename_list(),),
|
||||
'llama3_model_name':(s.get_filename_list(),),
|
||||
'text_encoder_model_name': (s.get_filename_list(),),
|
||||
'llm_vision_model_name':(s.get_filename_list(),),
|
||||
'type': base['required']['type'],
|
||||
'pretrained_path': ("STRING", {"multiline": False, "default": "HiTZ/Latxa-Llama-3.1-8B-Instruct"}),
|
||||
'instruct_from_pretrained': ("BOOLEAN", {"default": False, "tooltip": "Load Instruct model from pretrained_path"}),
|
||||
'encoder_pretrained_path': ("STRING", {"multiline": False, "default": "mlabonne/NeuralLlama-3-8B-Instruct-abliterated"}),
|
||||
'vision_pretrained_path': ("STRING", {"multiline": False, "default": "qresearch/llama-3.1-8B-vision-378"}),
|
||||
'encoder_from_pretrained': ("BOOLEAN", {"default": False, "tooltip": "Load Instruct model from pretrained_path"}),
|
||||
'vision_from_pretrained': ("BOOLEAN", {"default": False, "tooltip": "Load LLM+vision model from pretrained_path"}),
|
||||
'load_te': ("BOOLEAN", {"default": True, "tooltip": "If off, does not load separate model as text encoder, using only llm_vision_model_name"}),
|
||||
'system_prompt': ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "You are a creative assistant optimized for generating vivid, detailed descriptions for video generation."
|
||||
@@ -103,7 +106,7 @@ class FunPackDualCLIPLoaderInstruct:
|
||||
files += folder_paths.get_filename_list('clip')
|
||||
return sorted(files)
|
||||
|
||||
def load(self, clip_model_name, llama_instruct_model_name, llama3_model_name, type, system_prompt, pretrained_path, instruct_from_pretrained):
|
||||
def load(self, clip_model_name, type, text_encoder_model_name, llm_vision_model_name, encoder_pretrained_path, vision_pretrained_path, system_prompt, encoder_from_pretrained=None, vision_from_pretrained=None, load_te=None):
|
||||
# Load CLIP model using ComfyUI
|
||||
clip_path = folder_paths.get_full_path('clip', clip_model_name)
|
||||
def get_clip_type(type):
|
||||
@@ -111,36 +114,52 @@ class FunPackDualCLIPLoaderInstruct:
|
||||
print("Detected clip type:", clip_type)
|
||||
return clip_type
|
||||
|
||||
# Load LLaMA-3.1-Instruct from weights
|
||||
llama_path = folder_paths.get_full_path('clip', llama_instruct_model_name)
|
||||
config_source = pretrained_path
|
||||
|
||||
# Load regular LLaMA3 for compatibility
|
||||
llama3_path = folder_paths.get_full_path('clip', llama3_model_name)
|
||||
|
||||
# Load tokenizer and config
|
||||
tokenizer = AutoTokenizer.from_pretrained(config_source)
|
||||
config = LlamaConfig.from_pretrained(config_source)
|
||||
# Load TE model from weights
|
||||
encoder_path = folder_paths.get_full_path('clip', text_encoder_model_name)
|
||||
config_source = encoder_pretrained_path
|
||||
|
||||
if instruct_from_pretrained == True:
|
||||
model = LlamaForCausalLM.from_pretrained(config_source)
|
||||
pretrained_vision_saving_path = folder_paths.get_full_path('clip', vision_pretrained_path)
|
||||
|
||||
print("Loading TE from pretrained is set to", encoder_from_pretrained)
|
||||
print("Loading LLM+vision from pretrained is set to", vision_from_pretrained)
|
||||
print("Loading custom TE is set to", load_te)
|
||||
|
||||
# Load LLM with vision capabilities (expected llava-llama-3-8b_v1_1)
|
||||
if vision_from_pretrained == True:
|
||||
print("Loading LLM+vision from", vision_pretrained_path)
|
||||
model = AutoModelForCausalLM.from_pretrained(vision_pretrained_path, trust_remote_code=True)
|
||||
print("Saving model file to", pretrained_vision_saving_path)
|
||||
model.save_pretrained(f"{pretrained_vision_saving_path}")
|
||||
print("Pretrained model saved successfully.")
|
||||
vision_path = folder_paths.get_full_path('clip', vision_pretrained_path)
|
||||
else:
|
||||
model = LlamaForCausalLM(config)
|
||||
print("Loading LLM+vision from existing local safetensors file...")
|
||||
vision_path = folder_paths.get_full_path('clip', llm_vision_model_name)
|
||||
|
||||
try:
|
||||
print("Loading Instruct LLaMA from the path:", llama_path)
|
||||
state_dict = load_file(llama_path, device="cuda")
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
model.eval().to(torch.float16).requires_grad_(False)
|
||||
print("Llama Instruct model loaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"Error loading LLaMA Instruct model: {e}")
|
||||
raise
|
||||
if load_te == True:
|
||||
try:
|
||||
if encoder_from_pretrained == False:
|
||||
print("Loading custom text encoder from the path:", encoder_path)
|
||||
model = LlamaForCausalLM.from_pretrained(config_source, trust_remote_code=True)
|
||||
state_dict = load_file(encoder_path, device="cuda")
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
model.eval().to(torch.float16).requires_grad_(False)
|
||||
print("Custom text encoder from safetensors file loaded successfully!")
|
||||
else:
|
||||
print("Loading custom text encoder from the path:", encoder_pretrained_path)
|
||||
model = AutoModelForCausalLM.from_pretrained(encoder_pretrained_path, trust_remote_code=True)
|
||||
tokenizer = AutoTokenizer.from_pretrained(encoder_pretrained_path, trust_remote_code=True)
|
||||
model.load_state_dict(model, strict=False)
|
||||
model.eval().to(torch.float16).requires_grad_(False)
|
||||
print("Custom text encoder from transformers loaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"Error loading custom text encoder: {e}")
|
||||
raise
|
||||
|
||||
# Wrap it like a CLIP-compatible text encoder
|
||||
class InstructWrapper:
|
||||
def __init__(self):
|
||||
print("InstructWrapper initialized!")
|
||||
print("TEWrapper initialized!")
|
||||
self.system_prompt = system_prompt
|
||||
print("System prompt is set to:", self.system_prompt)
|
||||
def tokenize(self, text):
|
||||
@@ -158,18 +177,19 @@ class FunPackDualCLIPLoaderInstruct:
|
||||
return hidden, pooled
|
||||
|
||||
# Replace text encoder in CLIP model
|
||||
clip_model = sd.load_clip(ckpt_paths=[clip_path, llama3_path], embedding_directory=None, clip_type=get_clip_type(type), model_options={})
|
||||
clip_model.text = InstructWrapper()
|
||||
clip_model = sd.load_clip(ckpt_paths=[clip_path, vision_path], embedding_directory=None, clip_type=get_clip_type(type), model_options={})
|
||||
if load_te == True:
|
||||
clip_model.text = InstructWrapper()
|
||||
print("Current TE:", clip_model.text) # Check if encoder is replaced
|
||||
return (clip_model,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FunPackImg2LatentInterpolation": FunPackImg2LatentInterpolation,
|
||||
"FunPackDualCLIPLoaderInstruct": FunPackDualCLIPLoaderInstruct,
|
||||
"FunPackCLIPLoader": FunPackCLIPLoader,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FunPackImg2LatentInterpolation": "FunPack img2latent Interpolation",
|
||||
"FunPackDualCLIPLoaderInstruct": "FunPack DualCLIP Instruct Loader"
|
||||
"FunPackCLIPLoader": "FunPack CLIP Loader"
|
||||
}
|
||||
Reference in New Issue
Block a user