From 9a4629aba77fce0ec971ae7e032d40124c888baf Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sun, 16 Jun 2024 02:51:21 +0200
Subject: [PATCH] PixArt SD3 node prompt weight fix

#61
---
 PixArt/nodes.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/PixArt/nodes.py b/PixArt/nodes.py
index d9abda0..87ff74f 100644
--- a/PixArt/nodes.py
+++ b/PixArt/nodes.py
@@ -214,14 +214,26 @@ class PixArtT5FromSD3CLIP:
 
 	def split(self, sd3_clip, padding):
 		from comfy.sd3_clip import SD3Tokenizer, SD3ClipModel
+		import copy
 	
 		clip = sd3_clip.clone()
 		assert clip.cond_stage_model.t5xxl is not None, "CLIP must have T5 loaded!"
 
-		# model
+		# remove transformer
+		transformer = clip.cond_stage_model.t5xxl.transformer
+		clip.cond_stage_model.t5xxl.transformer = None
+
+		# clone object
 		tmp = SD3ClipModel(clip_l=False, clip_g=False, t5=False)
-		tmp.t5xxl = clip.cond_stage_model.t5xxl
-		
+		tmp.t5xxl = copy.deepcopy(clip.cond_stage_model.t5xxl)
+		# put transformer back
+		clip.cond_stage_model.t5xxl.transformer = transformer
+		tmp.t5xxl.transformer = transformer
+
+		# override special tokens
+		tmp.t5xxl.special_tokens = copy.deepcopy(clip.cond_stage_model.t5xxl.special_tokens)
+		tmp.t5xxl.special_tokens.pop("end") # make sure empty tokens match
+
 		# tokenizer
 		tok = SD3Tokenizer()
 		tok.t5xxl.min_length = padding