This commit is contained in:
DeepBeepMeep
2025-07-10 22:40:13 +02:00
parent 8d2164aaf1
commit 597d26b7e0
6 changed files with 234 additions and 102 deletions

View File

@@ -45,7 +45,7 @@ class DTT2V:
self.dtype = dtype
self.num_train_timesteps = config.num_train_timesteps
self.param_dtype = config.param_dtype
self.text_len = config.text_len
self.text_encoder = T5EncoderModel(
text_len=config.text_len,
dtype=config.t5_dtype,
@@ -250,11 +250,15 @@ class DTT2V:
if self._interrupt:
return None
text_len = self.text_len
prompt_embeds = self.text_encoder([input_prompt], self.device)[0]
prompt_embeds = prompt_embeds.to(self.dtype).to(self.device)
prompt_embeds = torch.cat([prompt_embeds, prompt_embeds.new_zeros(text_len -prompt_embeds.size(0), prompt_embeds.size(1)) ]).unsqueeze(0)
if self.do_classifier_free_guidance:
negative_prompt_embeds = self.text_encoder([n_prompt], self.device)[0]
negative_prompt_embeds = negative_prompt_embeds.to(self.dtype).to(self.device)
negative_prompt_embeds = torch.cat([negative_prompt_embeds, negative_prompt_embeds.new_zeros(text_len -negative_prompt_embeds.size(0), negative_prompt_embeds.size(1)) ]).unsqueeze(0)
if self._interrupt:
return None