From edf9f26e6d422491b206ef3db783d960a781b83e Mon Sep 17 00:00:00 2001 From: Olivv <133010426+olivv-cs@users.noreply.github.com> Date: Wed, 11 Jun 2025 18:45:23 +0300 Subject: [PATCH] Add files via upload --- funpack.py | 105 +++++++++++++++++++++++++++++++++++------------------ 1 file changed, 69 insertions(+), 36 deletions(-) diff --git a/funpack.py b/funpack.py index bc13d33d..7bee7987 100644 --- a/funpack.py +++ b/funpack.py @@ -198,10 +198,10 @@ class FunPackCLIPLoader: raise # Wrap it like a CLIP-compatible text encoder - class InstructWrapper: - def __init__(self, model, tokenizer, system_prompt, top_p, top_k, temperature, generate_assist_prompt=None): - print("TEWrapper initialized!") - self.model = model + 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 @@ -243,7 +243,16 @@ class FunPackCLIPLoader: {"role": "user", "content": text} ] - return self.tokenizer.apply_chat_template(messages, add_generation_prompt=False, return_tensors="pt") + 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()") @@ -252,11 +261,15 @@ class FunPackCLIPLoader: {"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" - ) + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=77 + ) with torch.no_grad(): generated_ids = self.model.generate( @@ -270,43 +283,63 @@ class FunPackCLIPLoader: ) 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): - with torch.no_grad(): - output = self.model(input_ids=tokens, output_hidden_states=True) - hidden_states = output.hidden_states[-1] - pooled_output = hidden_states.mean(dim=1) - # pooled = hidden[:, -1, :] # use last token for pooled - return {"l": pooled_output, "hidden_states": hidden_states} - + return self.clip_model.encode_from_tokens(tokens, return_pooled) + def encode_token_weights(self, tokens): - # This method is expected to return (hidden_states, pooled_output) - # where hidden_states is usually (batch_size, sequence_length, hidden_size) - # and pooled_output is (batch_size, hidden_size). - - # Since your encode_from_tokens already calculates these, we can reuse it. - # However, encode_from_tokens returns pooled_output directly, - # so we need to adjust to return both. - with torch.no_grad(): - output = self.model(input_ids=tokens, output_hidden_states=True) - hidden_states = output.hidden_states[-1] - pooled_output = hidden_states.mean(dim=1) # Or choose the last token's embedding for pooling - - return {"l": pooled_output, "hidden_states": hidden_states} - - def reset_clip_options(self): - return - + return self.clip_model.encode_token_weights(tokens) + + # Replace text encoder in CLIP model 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: - te = InstructWrapper(model, tokenizer, system_prompt, top_p, top_k, temperature) - #clip_model.tokenizer = te.tokenizer - #clip_model.cond_stage_model = model - #clip_model.tokenize = te.tokenize - clip_model.text = te.generate + + 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,) + else: return (clip_model,)