[HiDream] Hard-cap T5 embeddings to 128 length

This commit is contained in:
tdrussell 2025-04-16 13:42:36 -05:00
parent b4dc03ad76
commit bc388f8ffa

View File

@ -108,6 +108,8 @@ class HiDreamTEModel(torch.nn.Module):
if self.t5xxl is not None:
t5_output = self.t5xxl.encode_token_weights(token_weight_pairs_t5)
t5_out, t5_pooled = t5_output[:2]
# Quality rapidly degrades if T5 embeddings are >128 length. Llama3 is not as bad, so we allow that one to be longer.
t5_out = t5_out[:, :128, :]
if self.llama is not None:
ll_output = self.llama.encode_token_weights(token_weight_pairs_llama)