import os import torch from transformers import LlamaConfig, LlamaForCausalLM, LlavaForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer, AutoConfig from huggingface_hub import hf_hub_download, snapshot_download from safetensors.torch import load_file, save_file, safe_open import glob import comfy.model_management as mm import comfy.sd as sd import folder_paths import nodes import tempfile # import gc # No longer needed for aggressive unloading class FunPackImg2LatentInterpolation: @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE",), "frame_count": ("INT", {"default": 25, "min": 1, "max": 125, "step": 4}), } } RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("img_batch_for_encode", "img_for_start_image") FUNCTION = "process" CATEGORY = "WANPack/Video" def process(self, images, frame_count): device = images.device total_input_frames = images.shape[0] # 1. Take the last frame as our starting point last_frame = images[-1:] # 2. Generate smooth transition from this frame interpolated = [] # First frame is exact copy of last input frame (for perfect continuity) interpolated.append(last_frame.clone()) # Generate remaining frames with increasing denoise for i in range(1, frame_count): # Calculate denoise strength (0 at start, 1 at end) denoise_strength = i / (frame_count - 1) # Apply denoising noise = torch.randn_like(last_frame) blended = (1 - denoise_strength) * last_frame + denoise_strength * noise interpolated.append(blended) # Convert to tensor output = torch.cat(interpolated, dim=0) # Preview is first frame (same as input's last frame) preview = interpolated[0].clone() return (output, preview) # Helper to list files in models/clip 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 FunPackCLIPLoader: @classmethod def INPUT_TYPES(s): base = nodes.DualCLIPLoader.INPUT_TYPES() return { 'required': { 'clip_model_name': (s.get_filename_list(),), 'text_encoder_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"}), 'encoder_from_pretrained': ("BOOLEAN", {"default": False, "tooltip": "Load Instruct 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 an expert visual describer for AI video generation. Your task is to interpret user prompts and transform them into detailed, vivid descriptions optimized for image-to-video synthesis. Ensure your descriptions prioritize visual consistency, dynamic actions, and coherent scene elements to guide the generative model in creating smooth, logical video sequences from an initial image. Do not include conversational filler or explanations; just the descriptive text:<|eot_id|>" }), 'top_p': ("FLOAT", {"min": 0.0, "max": 10.0, "step": 0.05, "default": 0.75}), 'top_k': ("INT", {"min": 0, "max": 1000, "step": 1, "default": 40}), 'temperature': ("FLOAT", {"min": 0.0, "max": 10.0, "step": 0.01, "default": 0.6}), 'generate_assist_prompt': ("BOOLEAN", {"default": False}) } } RETURN_TYPES = 'CLIP', RETURN_NAMES = 'clip', FUNCTION = "load" CATEGORY = "conditioning" @classmethod def get_filename_list(s): files = [] 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, 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): clip_type = getattr(sd.CLIPType, type.upper(), sd.CLIPType.HUNYUAN_VIDEO) print("Detected clip type:", clip_type) return clip_type # Load TE model from weights encoder_path = folder_paths.get_full_path('clip', text_encoder_model_name) # 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 custom TE is set to", load_te) # 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: print("Loading custom text encoder from the path:", encoder_path) 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") # Load directly to GPU 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: print("Loading custom text encoder from the path:", encoder_pretrained_path) 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) 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 # 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, # Pass the loaded LLM model llm_tokenizer=llm_tokenizer, # Pass the loaded LLM tokenizer system_prompt=system_prompt, top_p=top_p, top_k=top_k, temperature=temperature, generate_assist_prompt=generate_assist_prompt ) # 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 # LLM model is now passed directly self.llm_tokenizer = llm_tokenizer # LLM tokenizer is now passed directly # Store the device where the LLM model is initially loaded # This is where the model will be moved to for inference self.llm_model_device = llm_model.device self.system_prompt = system_prompt self.top_p = top_p self.top_k = top_k self.temperature = temperature self.max_new_tokens = 1024 # Corrected token limit back to 1024 as per the original file 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 full detail.") print(f"LLM will be moved to/from: {self.llm_model_device}") 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}) # Move LLM to its designated device before generation self.llm_model.to(self.llm_model_device) print(f"[PromptEnhancerClipWrapper] LLM moved to {self.llm_model.device} for generation.") # This part of tokenization is for the LLM itself # Unpack the dictionary to get input_ids and attention_mask chat_template_output = self.llm_tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ) llm_input_ids = chat_template_output["input_ids"].to(self.llm_model.device) llm_attention_mask = chat_template_output["attention_mask"].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_input_ids, # Use the unpacked input_ids attention_mask=llm_attention_mask, # Pass the attention_mask 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_input_ids.shape[1]:], skip_special_tokens=True) self.assistant_reply = output_text # Store for next turn # Move LLM back to CPU after generation (reverted from full unload) self.llm_model.to("cpu") print(f"[PromptEnhancerClipWrapper] LLM moved back to CPU.") 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. # With max_new_tokens at 1024, the LLM will generate more detailed text. # CLIP's tokenizer will then handle its own internal truncation if the text is still too long, # but the truncated result should be of higher quality. 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, "FunPackCLIPLoader": FunPackCLIPLoader, } NODE_DISPLAY_NAME_MAPPINGS = { "FunPackImg2LatentInterpolation": "FunPack img2latent Interpolation", "FunPackCLIPLoader": "FunPack CLIP Loader" }