Revert an unnecessary change.

This commit is contained in:
Andrew Kvochko
2024-11-24 11:53:12 +00:00
parent 3b1f9cf1d3
commit ec9dedb5f0
+3 -3
View File
@@ -78,13 +78,13 @@ class LTXVTextEncoderModel(torch.nn.Module):
token_weight_pairs_t5 = token_weight_pairs["t5xxl"]
text_input_ids = torch.tensor(
[[t[0] for t in token_weight_pairs_t5]],
device=self.device,
device=self.t5xxl.device,
)
prompt_attention_mask = torch.tensor(
[[w[1] for w in token_weight_pairs_t5]],
device=self.device,
device=self.t5xxl.device,
)
self.to(self.device) # comfyui skips loading some weights to gpu
self.to(self.t5xxl.device) # comfyui skips loading some weights to gpu
out = self.t5xxl(text_input_ids, attention_mask=prompt_attention_mask)[0]
out = out * prompt_attention_mask.unsqueeze(2)
return out, None, {"attention_mask": prompt_attention_mask}