| 25 | |
| 26 | # Patch this to handle attention_mask, and also return attention_mask. |
| 27 | def encode_token_weights(self, token_weight_pairs): |
| 28 | token_weight_pairs = token_weight_pairs["gemma3_12b"] |
| 29 | |
| 30 | out, pooled, extra = self.gemma3_12b.encode_token_weights(token_weight_pairs) |
| 31 | out_device = out.device |
| 32 | if comfy.model_management.should_use_bf16(self.execution_device): |
| 33 | out = out.to(device=self.execution_device, dtype=torch.bfloat16) |
| 34 | |
| 35 | attention_mask = extra["attention_mask"] |
| 36 | bool_mask = attention_mask.bool().unsqueeze(1).unsqueeze(-1) |
| 37 | out = torch.where(bool_mask, out, torch.zeros_like(out)) |
| 38 | |
| 39 | assert self.text_projection_type == "dual_linear" |
| 40 | out = self.text_embedding_projection(out) |
| 41 | extra = {"unprocessed_ltxav_embeds": True, "attention_mask": attention_mask} |
| 42 | |
| 43 | return out.to(device=out_device, dtype=torch.float), pooled, extra |
| 44 | |
| 45 | comfy.text_encoders.lt.LTXAVTEModel.encode_token_weights = encode_token_weights |
| 46 | |