From ebb6192ae717b9f4d60fe2700e0f2d22a792abc9 Mon Sep 17 00:00:00 2001 From: Olivv <133010426+olivv-cs@users.noreply.github.com> Date: Wed, 11 Jun 2025 19:56:20 +0300 Subject: [PATCH] Complete refactor - thank you, Gemini! --- funpack.py | 330 ++++++++++++++++++++--------------------------------- 1 file changed, 122 insertions(+), 208 deletions(-) diff --git a/funpack.py b/funpack.py index 7bee7987..eb904cad 100644 --- a/funpack.py +++ b/funpack.py @@ -1,18 +1,3 @@ -""" -olivv's FunPack v1.1.2 - -Changelog: - -v1.1.2 - Added "instruct_from_pretrained" -v1.1.1 - Added DualCLIP Instruct loader node for experimenting with Instruct models in FramePack. -v1.1.0 - Changed interpolation logic. - -Custom ComfyUI node created specifically to imitate FramePack's approach, but using WAN2.1 models. -It takes the last frame of a video, then multiplies it many times, -gradually increasing denoising to the end, somewhat similar to what FramePack does, -but without frame prediction. - -""" import os import torch from transformers import LlamaConfig, LlamaForCausalLM, LlavaForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer, AutoConfig @@ -85,13 +70,10 @@ class FunPackCLIPLoader: { 'clip_model_name': (s.get_filename_list(),), 'text_encoder_model_name': (s.get_filename_list(),), - 'llm_vision_model_name':(s.get_filename_list(),), + 'llm_vision_model_name':(s.get_filename_list(),), # Retained for local loading 'type': base['required']['type'], 'encoder_pretrained_path': ("STRING", {"multiline": False, "default": "mlabonne/NeuralLlama-3-8B-Instruct-abliterated"}), - 'vision_pretrained_path': ("STRING", {"multiline": False, "default": "huihui-ai/Llama-3.2-11B-Vision-Instruct-abliterated"}), '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"}), - 'vision_from_pretrained_comfy': ("BOOLEAN", {"default": False, "tooltip": "Checks for the model in models/LLM instead of HuggingFace"}), '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, @@ -114,7 +96,7 @@ class FunPackCLIPLoader: files += folder_paths.get_filename_list('clip') return sorted(files) - def load(self, clip_model_name, type, text_encoder_model_name, llm_vision_model_name, encoder_pretrained_path, vision_pretrained_path, system_prompt, top_p, top_k, temperature, encoder_from_pretrained=None, vision_from_pretrained=None, vision_from_pretrained_comfy=None, load_te=None, generate_assist_prompt=None): + def load(self, clip_model_name, type, text_encoder_model_name, llm_vision_model_name, encoder_pretrained_path, system_prompt, top_p, top_k, temperature, encoder_from_pretrained=None, load_te=None, generate_assist_prompt=None): # Load CLIP model using ComfyUI clip_path = folder_paths.get_full_path('clip', clip_model_name) def get_clip_type(type): @@ -124,225 +106,157 @@ class FunPackCLIPLoader: # Load TE model from weights encoder_path = folder_paths.get_full_path('clip', text_encoder_model_name) - config_source = encoder_pretrained_path - if not vision_from_pretrained_comfy: - pretrained_vision_local_path = snapshot_download(repo_id=vision_pretrained_path) - pvlp_model = pretrained_vision_local_path + "/model.safetensors" - else: - pretrained_vision_local_path = folder_paths.models_dir + "/clip/" + vision_pretrained_path - print(pretrained_vision_local_path) - pvlp_model = pretrained_vision_local_path + "/model.safetensors" + # Simplified LLM+vision model loading: always load from existing local safetensors file + print("Loading LLM+vision from existing local safetensors file:", llm_vision_model_name) + vision_path = folder_paths.get_full_path('clip', llm_vision_model_name) 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: - print("Loading LLM+vision from", vision_pretrained_path) - if os.path.exists(pvlp_model): - print("Model already saved in a single file. Loading from local path...") - model_dir = snapshot_download(repo_id=vision_pretrained_path) - vision_path = pretrained_vision_local_path + "/model.safetensors" - print ("Loading from", vision_path) - else: - print("Local model does not exist. Loading, merging and saving it locally..") - if not vision_from_pretrained_comfy: - model_dir = snapshot_download(repo_id=vision_pretrained_path) - model = AutoModelForCausalLM.from_pretrained(vision_pretrained_path, ignore_mismatched_sizes=True, trust_remote_code=True) - else: - model_dir = folder_paths.models_dir + "/clip/" + vision_pretrained_path - shard_paths = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) - shard_paths = [p for p in shard_paths if "index" not in p] - print("Shard files:", shard_paths) - # Step 3: Load and combine shards - combined_state_dict = {} - for shard_path in shard_paths: - print(shard_path) - # Load shard to CPU to minimize memory usage - shard_state_dict = load_file(shard_path, device="cpu") - combined_state_dict.update(shard_state_dict) - # Free memory - del shard_state_dict - output_safetensors = os.path.join(pretrained_vision_local_path, "model.safetensors") - print("Model output path:", output_safetensors) - save_file(combined_state_dict, output_safetensors) - print("Shards successfully transformed into single model.") - vision_path = pretrained_vision_local_path + "/model.safetensors" - else: - print("Loading LLM+vision from existing local safetensors file...") - vision_path = folder_paths.get_full_path('clip', llm_vision_model_name) - + # Initialize tokenizer outside the PromptEnhancer for the LLM + llm_tokenizer = None + te = None # This will hold the LLM model if load_te: try: if not encoder_from_pretrained: - te = None print("Loading custom text encoder from the path:", encoder_path) - tokenizer = AutoTokenizer.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True) + llm_tokenizer = AutoTokenizer.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True) model = LlamaForCausalLM.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True) state_dict = load_file(encoder_path, device="cuda") model.load_state_dict(state_dict, strict=False) te = model.eval().to(torch.float16).requires_grad_(False) print("Custom text encoder from safetensors file loaded successfully!") else: - te = None print("Loading custom text encoder from the path:", encoder_pretrained_path) - config = AutoConfig.from_pretrained(encoder_pretrained_path, trust_remote_code=True) + llm_tokenizer = AutoTokenizer.from_pretrained(encoder_pretrained_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(encoder_pretrained_path, ignore_mismatched_sizes=True, trust_remote_code=True) - tokenizer = AutoTokenizer.from_pretrained(encoder_pretrained_path, trust_remote_code=True) te = 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 PromptEnhancer: - def __init__(self, te, tokenizer, system_prompt, top_p, top_k, temperature, generate_assist_prompt=None): - print("Prompt Enhancer initialized!") - self.model = te - self.tokenizer = tokenizer - self.system_prompt = system_prompt - self.top_p = top_p - self.top_k = top_k - self.temperature = temperature - self.max_new_tokens = 256 - self.assistant_reply = None - self.assistant_replies = {} - self.generate_assist_prompt = generate_assist_prompt - print("top_p:", self.top_p) - print("top_k:", self.top_k) - print("temperature:", self.temperature) - print("max new tokens is hardcoded to", self.max_new_tokens) - print("System prompt is set to:", self.system_prompt) - - def tokenize(self, text): - print("Calling tokenize()", flush=True) - if text not in self.assistant_replies and self.generate_assist_prompt: - print("Generating assistant reply..") - self.assistant_reply = self.generate(text) - self.assistant_replies[text] = self.assistant_reply - print("Assistant generated:", self.assistant_reply, flush=True) - elif text in self.assistant_replies and self.generate_assist_prompt: - print("Using cached assistant reply.") - self.assistant_reply = self.assistant_replies[text] - else: - print("Assistant prompt has been disabled.") - - if self.generate_assist_prompt: - messages = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": text}, - {"role": "assistant", "content": self.assistant_reply} - ] - - else: - messages = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": text} - ] - - tokens = self.tokenizer.apply_chat_template( - messages, - add_generation_prompt=False, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=77 - ) - - return tokens - - def generate(self, text): - print("Calling generate() from tokenize()") - # Prepare chat messages for assistant generation - messages = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": text} - ] - - tokens = self.tokenizer.apply_chat_template( - messages, - add_generation_prompt=True, # Important for generation - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=77 - ) - - with torch.no_grad(): - generated_ids = self.model.generate( - input_ids=tokens, - do_sample=True, - top_p=self.top_p, - top_k=self.top_k, - temperature=self.temperature, - max_new_tokens=self.max_new_tokens, - pad_token_id=tokenizer.pad_token_id - ) - output_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True) - return output_text - - class PromptRewriter: - - def __init__(self, enhancer, clip_model): - self.enhancer = enhancer - self.clip_model = clip_model - self.tokenizer = clip_model.tokenizer - self.ep = None - self.enhanced_prompts = {} - - def tokenize(self, prompt): - # Let your LLM rewrite the prompt - - if self.ep not in self.enhanced_prompts: - print("Generating enhanced prompt..") - self.enhanced_prompt = self.enhancer.generate(prompt) - self.enhanced_prompts[self.ep] = self.enhanced_prompt - print("Assistant generated:", self.enhanced_prompt, flush=True) - elif self.ep in self.enhanced_prompts: - print("Using cached enhanced prompt.") - self.enhanced_prompt = self.enhanced_prompts[self.ep] - else: - print("Enhanced prompt is broken or somehow has been disabled.") - - print(f"[PromptRewriter] Rewritten prompt: {self.enhanced_prompt}") - # Now pass it to CLIP tokenizer - return self.clip_model.tokenize(enhanced_prompt) - - def encode_from_tokens(self, tokens, return_pooled=True): - return self.clip_model.encode_from_tokens(tokens, return_pooled) - - def encode_token_weights(self, tokens): - return self.clip_model.encode_token_weights(tokens) - - - # Replace text encoder in CLIP model + # Load the base CLIP model first clip_model = sd.load_clip(ckpt_paths=[clip_path, vision_path], embedding_directory=None, clip_type=get_clip_type(type), model_options={"ignore_mismatched_sizes": True}) if load_te == True: + # Wrap the original CLIP model with our prompt enhancer + wrapped_clip_with_enhancement = PromptEnhancerClipWrapper( + original_clip_model=clip_model, + llm_model=te, + llm_tokenizer=llm_tokenizer, + system_prompt=system_prompt, + top_p=top_p, + top_k=top_k, + temperature=temperature, + generate_assist_prompt=generate_assist_prompt + ) - enhancer = PromptEnhancer(te, tokenizer, system_prompt, top_p, top_k, temperature) - - wrapped_clip = PromptRewriter(enhancer, clip_model) - - clip_model.tokenize = wrapped_clip.tokenize - - clip_model.tokenizer = wrapped_clip.tokenizer - - clip_model.encode_from_tokens = wrapped_clip.encode_from_tokens - - clip_model.encode_token_weights = wrapped_clip.encode_token_weights - - print("Current TE:", clip_model.cond_stage_model) # Check if encoder is replaced - - - return (clip_model,) + # Return our wrapper instance directly. ComfyUI will then call methods on it. + return (wrapped_clip_with_enhancement,) else: + # If no TE, return the original CLIP model without enhancement return (clip_model,) +class PromptEnhancerClipWrapper: + def __init__(self, original_clip_model, llm_model, llm_tokenizer, system_prompt, top_p, top_k, temperature, generate_assist_prompt): + self.original_clip = original_clip_model # Keep a reference to the actual CLIP object + self.llm_model = llm_model + self.llm_tokenizer = llm_tokenizer + self.system_prompt = system_prompt + self.top_p = top_p + self.top_k = top_k + self.temperature = temperature + self.max_new_tokens = 77 # Adjusted to try and fit CLIP's typical 77-token limit + self.assistant_reply = None # For conversational context if generate_assist_prompt is true + self.generate_assist_prompt = generate_assist_prompt + + self.enhanced_prompts_cache = {} # Stores original_user_prompt -> enhanced_text + self._is_llm_generating = False # Flag to prevent recursive LLM calls if LLM's own tokenization process calls us + + print("[PromptEnhancerClipWrapper] Initialized with LLM and CLIP.") + print(f"top_p: {self.top_p}, top_k: {self.top_k}, temperature: {self.temperature}") + print(f"System prompt: {self.system_prompt}") + print(f"LLM max_new_tokens set to: {self.max_new_tokens} for CLIP compatibility.") + + def _generate_enhanced_prompt(self, user_prompt): + print("[PromptEnhancerClipWrapper] Calling _generate_enhanced_prompt for LLM inference...") + messages = [ + {"role": "system", "content": self.system_prompt}, + {"role": "user", "content": user_prompt} + ] + + # If generate_assist_prompt is true and we have a previous assistant reply, add it + if self.generate_assist_prompt and self.assistant_reply: + messages.append({"role": "assistant", "content": self.assistant_reply}) + + # This part of tokenization is for the LLM itself + llm_tokens = self.llm_tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + return_tensors="pt", + ).to(self.llm_model.device) + + with torch.no_grad(): + self._is_llm_generating = True # Set flag before LLM generation + try: + generated_ids = self.llm_model.generate( + input_ids=llm_tokens, + do_sample=True, + top_p=self.top_p, + top_k=self.top_k, + temperature=self.temperature, + max_new_tokens=self.max_new_tokens, + pad_token_id=self.llm_tokenizer.pad_token_id + ) + finally: + self._is_llm_generating = False # Always reset flag + + # Decode only the newly generated part + output_text = self.llm_tokenizer.decode(generated_ids[0][llm_tokens.shape[1]:], skip_special_tokens=True) + self.assistant_reply = output_text # Store for next turn + return output_text + + def tokenize(self, prompt): + print(f"[PromptEnhancerClipWrapper] Received prompt for tokenization: '{prompt}'") + + # If LLM is currently generating, this call is likely from LLM's internal tokenization (e.g. from tokenizer.apply_chat_template). + # Pass through to original CLIP tokenizer directly. + if self._is_llm_generating: + print(f"[PromptEnhancerClipWrapper] LLM generation in progress. Passing directly to original CLIP tokenizer: '{prompt}'") + return self.original_clip.tokenize(prompt) + + # Check if the prompt is an already enhanced prompt (a value in our cache) + # This prevents re-enhancing text that was already LLM-generated. + if prompt in self.enhanced_prompts_cache.values(): + print(f"[PromptEnhancerClipWrapper] Prompt '{prompt}' is an already enhanced prompt. Passing directly to original CLIP tokenizer.") + return self.original_clip.tokenize(prompt) + + # If it's a new, original user prompt (key not in cache) + if prompt not in self.enhanced_prompts_cache: + print(f"[PromptEnhancerClipWrapper] Prompt '{prompt}' NOT in cache. Calling LLM for enhancement...") + enhanced_text = self._generate_enhanced_prompt(prompt) + self.enhanced_prompts_cache[prompt] = enhanced_text + print(f"[PromptEnhancerClipWrapper] LLM enhanced and cached for '{prompt}': {enhanced_text}") + else: + print(f"[PromptEnhancerClipWrapper] Using cached enhanced prompt for '{prompt}'.") + enhanced_text = self.enhanced_prompts_cache[prompt] + + print(f"[PromptEnhancerClipWrapper] Final prompt being passed to CLIP: {enhanced_text}") + + # Pass the enhanced text to the original CLIP's tokenize method. + # The 'max_new_tokens' setting on the LLM should now ensure the string length is suitable. + return self.original_clip.tokenize(enhanced_text) + + # These methods simply delegate to the original CLIP model + def encode_from_tokens(self, tokens, return_pooled=True): + return self.original_clip.encode_from_tokens(tokens, return_pooled) + + def encode_token_weights(self, tokens): + return self.original_clip.encode_token_weights(tokens) + NODE_CLASS_MAPPINGS = { "FunPackImg2LatentInterpolation": FunPackImg2LatentInterpolation,