nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Token limit = 512 + loading optimizations

Browse Source
This commit is contained in:
Olivv
2025-06-12 00:35:20 +03:00
committed by GitHub
parent 8721dc4048
commit a7d277438b
1 changed files with 58 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
funpack.py
+58 -31
View File
@@ -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}'")