Update modeling_latte.py

This commit is contained in:
lb203
2024-05-27 22:57:34 +08:00
committed by GitHub
parent dc12340706
commit b5b2c0968e
@@ -401,7 +401,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
encoder_hidden_states = torch.cat([encoder_hidden_states_video, encoder_hidden_states_image], dim=1)
encoder_hidden_states_spatial = rearrange(encoder_hidden_states, 'b f t d -> (b f) t d').contiguous()
else:
encoder_hidden_states_spatial = repeat(encoder_hidden_states, 'b t d -> (b f) t d', f=frame).contiguous()
encoder_hidden_states_spatial = repeat(encoder_hidden_states, 'b 1 t d -> (b f) t d', f=frame).contiguous()
# prepare timesteps for spatial and temporal block
timestep_spatial = repeat(timestep, 'b d -> (b f) d', f=frame + use_image_num).contiguous()
@@ -673,4 +673,4 @@ if __name__ == '__main__':
encoder_attention_mask=cond_mask, use_image_num=args.use_image_num, timestep=timestep)
with torch.no_grad():
output = model(**model_kwargs)
# print(output)
# print(output)