From 45afdadf420d184b6764e82dcec149d8526ffa28 Mon Sep 17 00:00:00 2001
From: Olivv <133010426+olivv-cs@users.noreply.github.com>
Date: Wed, 11 Jun 2025 13:41:21 +0300
Subject: [PATCH]  Expected all tensors to be on the same device

---
 funpack.py | 31 ++++++++++++-------------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/funpack.py b/funpack.py
index c65af271..dcd09382 100644
--- a/funpack.py
+++ b/funpack.py
@@ -183,25 +183,23 @@ class FunPackCLIPLoader:
                     tokenizer = AutoTokenizer.from_pretrained(config_source, trust_remote_code=True)
                     state_dict = load_file(encoder_path, device="cuda")
                     model.load_state_dict(state_dict, strict=False)
-                    model.eval().to(torch.float16).requires_grad_(False)
+                    te = model.eval().to(torch.float16).requires_grad_(False)
                     print("Custom text encoder from safetensors file loaded successfully!")
-                    return te
                 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)
                     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)
-                    model.eval().to(torch.float16).requires_grad_(False)
+                    te = model.eval().to(torch.float16).requires_grad_(False)
                     print("Custom text encoder from transformers loaded successfully!")
-                    return te
             except Exception as e:
                 print(f"Error loading custom text encoder: {e}")
                 raise
 
         # Wrap it like a CLIP-compatible text encoder
         class InstructWrapper:
-            def __init__(self, model, tokenizer):
+            def __init__(self, model, tokenizer, system_prompt, top_p, top_k, temperature):
                 print("TEWrapper initialized!")
                 self.model = model
                 self.tokenizer = tokenizer
@@ -249,32 +247,27 @@ class FunPackCLIPLoader:
                         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)
+                    output_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True).to("cuda")
                     return output_text
 
             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_output = hidden_states.mean(dim=1).to("cuda")
                     # pooled = hidden[:, -1, :]  # use last token for pooled
                 return pooled_output
 
         # 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:
-            text_encoder = sd.load_text_encoder_state_dicts(state_dicts = [te], embedding_directory=None, clip_type = get_clip_type(type), model_options={})
-            #clip_model.text = InstructWrapper(
-            #    model=model,
-            #    tokenizer=tokenizer,
-            #    system_prompt=system_prompt,
-            #    top_p=top_p,
-            #    top_k=top_k,
-            #    temperature=temperature
-            #)
-            #print("Current TE:", clip_model.text)  # Check if encoder is replaced
-            print("Current TE:", text_encoder)
-            return (clip_model, text_encoder,)
+            te = InstructWrapper(model, tokenizer, system_prompt, top_p, top_k, temperature)
+            clip_model.cond_stage_model = te
+            clip_model.tokenizer = te.tokenizer
+            clip_model.tokenize = te.tokenize
+            print("Current TE:", clip_model.cond_stage_model)  # Check if encoder is replaced
+            print("clip_model attributes:", dir(clip_model))
+            return (clip_model,)
         else:
             return (clip_model,)