@@ -79,27 +79,41 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
7979 (_ , tokenizer ) = exit_stack .enter_context (tokenizer_info .model_on_device ())
8080
8181 # Apply LoRA models to the text encoder
82+ lora_dtype = TorchDevice .choose_bfloat16_safe_dtype (device )
8283 exit_stack .enter_context (
8384 LayerPatcher .apply_smart_model_patches (
8485 model = text_encoder ,
8586 patches = self ._lora_iterator (context ),
8687 prefix = Z_IMAGE_LORA_QWEN3_PREFIX ,
87- dtype = torch . bfloat16 ,
88+ dtype = lora_dtype ,
8889 )
8990 )
9091
9192 context .util .signal_progress ("Running Qwen3 text encoder" )
92- assert isinstance (text_encoder , PreTrainedModel )
93- assert isinstance (tokenizer , PreTrainedTokenizerBase )
93+ if not isinstance (text_encoder , PreTrainedModel ):
94+ raise TypeError (
95+ f"Expected PreTrainedModel for text encoder, got { type (text_encoder ).__name__ }
96+ "The Qwen3 encoder model may be corrupted or incompatible."
97+ )
98+ if not isinstance (tokenizer , PreTrainedTokenizerBase ):
99+ raise TypeError (
100+ f"Expected PreTrainedTokenizerBase for tokenizer, got { type (tokenizer ).__name__ }
101+ "The Qwen3 tokenizer may be corrupted or incompatible."
102+ )
94103
95104 # Apply chat template similar to diffusers ZImagePipeline
96105 # The chat template formats the prompt for the Qwen3 model
97- prompt_formatted = tokenizer .apply_chat_template (
98- [{"role" : "user" , "content" : prompt }],
99- tokenize = False ,
100- add_generation_prompt = True ,
101- enable_thinking = True ,
102- )
106+ try :
107+ prompt_formatted = tokenizer .apply_chat_template (
108+ [{"role" : "user" , "content" : prompt }],
109+ tokenize = False ,
110+ add_generation_prompt = True ,
111+ enable_thinking = True ,
112+ )
113+ except (AttributeError , TypeError ) as e :
114+ # Fallback if tokenizer doesn't support apply_chat_template or enable_thinking
115+ context .logger .warning (f"Chat template failed ({ e } )
116+ prompt_formatted = prompt
103117
104118 # Tokenize the formatted prompt
105119 text_inputs = tokenizer (
@@ -113,8 +127,16 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
113127
114128 text_input_ids = text_inputs .input_ids
115129 attention_mask = text_inputs .attention_mask
116- assert isinstance (text_input_ids , torch .Tensor )
117- assert isinstance (attention_mask , torch .Tensor )
130+ if not isinstance (text_input_ids , torch .Tensor ):
131+ raise TypeError (
132+ f"Expected torch.Tensor for input_ids, got { type (text_input_ids ).__name__ }
133+ "Tokenizer returned unexpected type."
134+ )
135+ if not isinstance (attention_mask , torch .Tensor ):
136+ raise TypeError (
137+ f"Expected torch.Tensor for attention_mask, got { type (attention_mask ).__name__ }
138+ "Tokenizer returned unexpected type."
139+ )
118140
119141 # Check for truncation
120142 untruncated_ids = tokenizer (prompt_formatted , padding = "longest" , return_tensors = "pt" ).input_ids
@@ -135,6 +157,18 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
135157 attention_mask = prompt_mask ,
136158 output_hidden_states = True ,
137159 )
160+
161+ # Validate hidden_states output
162+ if not hasattr (outputs , "hidden_states" ) or outputs .hidden_states is None :
163+ raise RuntimeError (
164+ "Text encoder did not return hidden_states. "
165+ "Ensure output_hidden_states=True is supported by this model."
166+ )
167+ if len (outputs .hidden_states ) < 2 :
168+ raise RuntimeError (
169+ f"Expected at least 2 hidden states from text encoder, got { len (outputs .hidden_states )}
170+ "This may indicate an incompatible model or configuration."
171+ )
138172 prompt_embeds = outputs .hidden_states [- 2 ]
139173
140174 # Z-Image expects a 2D tensor [seq_len, hidden_dim] with only valid tokens
@@ -143,13 +177,21 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
143177 # Since batch_size=1, we take the first item and filter by mask
144178 prompt_embeds = prompt_embeds [0 ][prompt_mask [0 ]]
145179
146- assert isinstance (prompt_embeds , torch .Tensor )
180+ if not isinstance (prompt_embeds , torch .Tensor ):
181+ raise TypeError (
182+ f"Expected torch.Tensor for prompt embeddings, got { type (prompt_embeds ).__name__ }
183+ "Text encoder returned unexpected type."
184+ )
147185 return prompt_embeds
148186
149187 def _lora_iterator (self , context : InvocationContext ) -> Iterator [Tuple [ModelPatchRaw , float ]]:
150188 """Iterate over LoRA models to apply to the Qwen3 text encoder."""
151189 for lora in self .qwen3_encoder .loras :
152190 lora_info = context .models .load (lora .lora )
153- assert isinstance (lora_info .model , ModelPatchRaw )
191+ if not isinstance (lora_info .model , ModelPatchRaw ):
192+ raise TypeError (
193+ f"Expected ModelPatchRaw for LoRA '{ lora .lora .key } { type (lora_info .model ).__name__ }
194+ "The LoRA model may be corrupted or incompatible."
195+ )
154196 yield (lora_info .model , lora .weight )
155197 del lora_info
0 commit comments