@@ -63,27 +63,41 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
6363 (_ , tokenizer ) = exit_stack .enter_context (tokenizer_info .model_on_device ())
6464
6565 # Apply LoRA models to the text encoder
66+ lora_dtype = TorchDevice .choose_bfloat16_safe_dtype (device )
6667 exit_stack .enter_context (
6768 LayerPatcher .apply_smart_model_patches (
6869 model = text_encoder ,
6970 patches = self ._lora_iterator (context ),
7071 prefix = Z_IMAGE_LORA_QWEN3_PREFIX ,
71- dtype = torch . bfloat16 ,
72+ dtype = lora_dtype ,
7273 )
7374 )
7475
7576 context .util .signal_progress ("Running Qwen3 text encoder" )
76- assert isinstance (text_encoder , PreTrainedModel )
77- assert isinstance (tokenizer , PreTrainedTokenizerBase )
77+ if not isinstance (text_encoder , PreTrainedModel ):
78+ raise TypeError (
79+ f"Expected PreTrainedModel for text encoder, got { type (text_encoder ).__name__ }
80+ "The Qwen3 encoder model may be corrupted or incompatible."
81+ )
82+ if not isinstance (tokenizer , PreTrainedTokenizerBase ):
83+ raise TypeError (
84+ f"Expected PreTrainedTokenizerBase for tokenizer, got { type (tokenizer ).__name__ }
85+ "The Qwen3 tokenizer may be corrupted or incompatible."
86+ )
7887
7988 # Apply chat template similar to diffusers ZImagePipeline
8089 # The chat template formats the prompt for the Qwen3 model
81- prompt_formatted = tokenizer .apply_chat_template (
82- [{"role" : "user" , "content" : prompt }],
83- tokenize = False ,
84- add_generation_prompt = True ,
85- enable_thinking = True ,
86- )
90+ try :
91+ prompt_formatted = tokenizer .apply_chat_template (
92+ [{"role" : "user" , "content" : prompt }],
93+ tokenize = False ,
94+ add_generation_prompt = True ,
95+ enable_thinking = True ,
96+ )
97+ except (AttributeError , TypeError ) as e :
98+ # Fallback if tokenizer doesn't support apply_chat_template or enable_thinking
99+ context .logger .warning (f"Chat template failed ({ e } )
100+ prompt_formatted = prompt
87101
88102 # Tokenize the formatted prompt
89103 text_inputs = tokenizer (
@@ -97,8 +111,16 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
97111
98112 text_input_ids = text_inputs .input_ids
99113 attention_mask = text_inputs .attention_mask
100- assert isinstance (text_input_ids , torch .Tensor )
101- assert isinstance (attention_mask , torch .Tensor )
114+ if not isinstance (text_input_ids , torch .Tensor ):
115+ raise TypeError (
116+ f"Expected torch.Tensor for input_ids, got { type (text_input_ids ).__name__ }
117+ "Tokenizer returned unexpected type."
118+ )
119+ if not isinstance (attention_mask , torch .Tensor ):
120+ raise TypeError (
121+ f"Expected torch.Tensor for attention_mask, got { type (attention_mask ).__name__ }
122+ "Tokenizer returned unexpected type."
123+ )
102124
103125 # Check for truncation
104126 untruncated_ids = tokenizer (prompt_formatted , padding = "longest" , return_tensors = "pt" ).input_ids
@@ -119,6 +141,18 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
119141 attention_mask = prompt_mask ,
120142 output_hidden_states = True ,
121143 )
144+
145+ # Validate hidden_states output
146+ if not hasattr (outputs , "hidden_states" ) or outputs .hidden_states is None :
147+ raise RuntimeError (
148+ "Text encoder did not return hidden_states. "
149+ "Ensure output_hidden_states=True is supported by this model."
150+ )
151+ if len (outputs .hidden_states ) < 2 :
152+ raise RuntimeError (
153+ f"Expected at least 2 hidden states from text encoder, got { len (outputs .hidden_states )}
154+ "This may indicate an incompatible model or configuration."
155+ )
122156 prompt_embeds = outputs .hidden_states [- 2 ]
123157
124158 # Z-Image expects a 2D tensor [seq_len, hidden_dim] with only valid tokens
@@ -127,13 +161,21 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.
127161 # Since batch_size=1, we take the first item and filter by mask
128162 prompt_embeds = prompt_embeds [0 ][prompt_mask [0 ]]
129163
130- assert isinstance (prompt_embeds , torch .Tensor )
164+ if not isinstance (prompt_embeds , torch .Tensor ):
165+ raise TypeError (
166+ f"Expected torch.Tensor for prompt embeddings, got { type (prompt_embeds ).__name__ }
167+ "Text encoder returned unexpected type."
168+ )
131169 return prompt_embeds
132170
133171 def _lora_iterator (self , context : InvocationContext ) -> Iterator [Tuple [ModelPatchRaw , float ]]:
134172 """Iterate over LoRA models to apply to the Qwen3 text encoder."""
135173 for lora in self .qwen3_encoder .loras :
136174 lora_info = context .models .load (lora .lora )
137- assert isinstance (lora_info .model , ModelPatchRaw )
175+ if not isinstance (lora_info .model , ModelPatchRaw ):
176+ raise TypeError (
177+ f"Expected ModelPatchRaw for LoRA '{ lora .lora .key } { type (lora_info .model ).__name__ }
178+ "The LoRA model may be corrupted or incompatible."
179+ )
138180 yield (lora_info .model , lora .weight )
139181 del lora_info
0 commit comments