From a7d277438b56a92ac70277ec0fabff3e95ef070b Mon Sep 17 00:00:00 2001 From: Olivv <133010426+olivv-cs@users.noreply.github.com> Date: Thu, 12 Jun 2025 00:35:20 +0300 Subject: [PATCH] Token limit = 512 + loading optimizations --- funpack.py | 89 +++++++++++++++++++++++++++++++++++------------------- 1 file changed, 58 insertions(+), 31 deletions(-) diff --git a/funpack.py b/funpack.py index 159acdef..6d32e15d 100644 --- a/funpack.py +++ b/funpack.py @@ -9,6 +9,7 @@ import comfy.sd as sd import folder_paths import nodes import tempfile +import gc # Import garbage collector class FunPackImg2LatentInterpolation: @classmethod @@ -114,38 +115,17 @@ class FunPackCLIPLoader: 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") - 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 + # Pass loading parameters instead of actual model/tokenizer instances wrapped_clip_with_enhancement = PromptEnhancerClipWrapper( original_clip_model=clip_model, - llm_model=te, - llm_tokenizer=llm_tokenizer, + encoder_path=encoder_path, + encoder_pretrained_path=encoder_pretrained_path, + encoder_from_pretrained=encoder_from_pretrained, system_prompt=system_prompt, top_p=top_p, top_k=top_k, @@ -161,10 +141,18 @@ class FunPackCLIPLoader: 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): + def __init__(self, original_clip_model, encoder_path, encoder_pretrained_path, encoder_from_pretrained, 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 + + # Store loading parameters for on-demand loading + self.encoder_path = encoder_path + self.encoder_pretrained_path = encoder_pretrained_path + self.encoder_from_pretrained = encoder_from_pretrained + + self.llm_model = None # Will be loaded on demand + self.llm_tokenizer = None # Will be loaded on demand + self.llm_model_device = None # Will be set upon loading + self.system_prompt = system_prompt self.top_p = top_p self.top_k = top_k @@ -176,7 +164,7 @@ class PromptEnhancerClipWrapper: 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("[PromptEnhancerClipWrapper] Initialized for on-demand LLM loading/unloading.") 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.") @@ -192,7 +180,31 @@ class PromptEnhancerClipWrapper: 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 + # Load LLM model and tokenizer only if not already loaded (e.g., first call in a generation cycle) + if self.llm_model is None or self.llm_tokenizer is None: + print("[PromptEnhancerClipWrapper] LLM model and tokenizer not loaded. Loading now...") + try: + if not self.encoder_from_pretrained: + self.llm_tokenizer = AutoTokenizer.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True) + model_base = LlamaForCausalLM.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True) + state_dict = load_file(self.encoder_path, device="cpu") # Load to CPU initially to control device + model_base.load_state_dict(state_dict, strict=False) + self.llm_model = model_base.eval().to(torch.float16).requires_grad_(False) # Convert to float16 on CPU + print(f"Custom text encoder from safetensors file loaded successfully to CPU!") + else: + self.llm_tokenizer = AutoTokenizer.from_pretrained(self.encoder_pretrained_path, trust_remote_code=True) + self.llm_model = AutoModelForCausalLM.from_pretrained(self.encoder_pretrained_path, ignore_mismatched_sizes=True, trust_remote_code=True) + self.llm_model = self.llm_model.eval().to(torch.float16).requires_grad_(False) # Convert to float16 on CPU + print(f"Custom text encoder from transformers loaded successfully to CPU!") + + self.llm_model_device = mm.get_torch_device() # Get the actual GPU device + self.llm_model.to(self.llm_model_device) # Move to GPU for inference + print(f"[PromptEnhancerClipWrapper] LLM moved to {self.llm_model.device} for generation.") + + except Exception as e: + print(f"Error loading LLM model for enhancement: {e}") + raise + llm_tokens = self.llm_tokenizer.apply_chat_template( messages, add_generation_prompt=True, @@ -217,7 +229,22 @@ class PromptEnhancerClipWrapper: # 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 + + # Fully unload LLM model and tokenizer after use + print("[PromptEnhancerClipWrapper] Attempting to fully unload LLM model...") + if self.llm_model: + self.llm_model.cpu() # Ensure it's on CPU before deletion + del self.llm_model + self.llm_model = None + if self.llm_tokenizer: + del self.llm_tokenizer + self.llm_tokenizer = None + + torch.cuda.empty_cache() # Clear GPU cache + gc.collect() # Force Python garbage collection + print(f"[PromptEnhancerClipWrapper] LLM model fully unloaded (from GPU and attempting to free RAM).") + + return output_text def tokenize(self, prompt): print(f"[PromptEnhancerClipWrapper] Received prompt for tokenization: '{prompt}'")