11from comfy import sd1_clip
22import comfy .text_encoders .llama
3- from .qwen_image import QwenImageTokenizer , QwenImageTEModel
3+ from .qwen_image import QwenImageTokenizer
44from transformers import ByT5Tokenizer
55import os
66import re
@@ -117,6 +117,8 @@ def load_sd(self, sd):
117117 return self .byt5_small .load_sd (sd )
118118 else :
119119 return super ().load_sd (sd )
120+
121+
120122class HunyuanImageRefinerTEModel (sd1_clip .SD1ClipModel ):
121123 def __init__ (self , device = "cpu" , dtype = None , model_options = {}):
122124 super ().__init__ (device = device , dtype = dtype , name = "qwen25_7b" , clip_model = Qwen25_7BVLIModel , model_options = model_options )
@@ -142,6 +144,7 @@ def encode_token_weights(self, token_weight_pairs):
142144
143145
144146def te (byt5 = True , dtype_llama = None , llama_scaled_fp8 = None , refiner = False ):
147+
145148 class HunyuanImageTEModel_ (HunyuanImageTEModel ):
146149
147150 def __init__ (self , device = "cpu" , dtype = None , model_options = {}):
@@ -151,6 +154,7 @@ def __init__(self, device="cpu", dtype=None, model_options={}):
151154 if dtype_llama is not None :
152155 dtype = dtype_llama
153156 super ().__init__ (byt5 = byt5 , device = device , dtype = dtype , model_options = model_options )
157+
154158 class HunyuanImageTEModel_refiner (HunyuanImageRefinerTEModel ):
155159 def __init__ (self , device = "cpu" , dtype = None , model_options = {}):
156160 if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options :
@@ -161,4 +165,5 @@ def __init__(self, device="cpu", dtype=None, model_options={}):
161165 assert refiner , "refiner must be True"
162166 assert not byt5 , "byt5 must be False"
163167 super ().__init__ (device = device , dtype = dtype , model_options = model_options )
168+
164169 return HunyuanImageTEModel_refiner if refiner else HunyuanImageTEModel_
0 commit comments