From 11996dc9b4e09d20a6b493af3f7c2bb75a6409fa Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Thu, 7 Mar 2024 16:41:31 +0800 Subject: [PATCH] fixed dit forward Former-commit-id: 4cd537fab12231f7471107762b02b36217ca42d8 --- opensora/models/ae/videobase/vqvae/videogpt/rec_video.py | 5 ++--- opensora/models/diffusion/dit/dit.py | 5 +++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py index 5bc2d2f19..0353334d2 100644 --- a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py +++ b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py @@ -89,9 +89,8 @@ def main(args): x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) x_vae = x_vae.to(device) - with torch.no_grad(): - encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) - video_recon = vqvae.decode(encodings) + encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) + video_recon = vqvae.decode(encodings) # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4') custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path) diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index d7611b2b5..2c5e33348 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -313,7 +313,7 @@ class DiT(nn.Module): attention_mask = attention_mask.masked_fill(attention_mask == 0, 1e-8) return attention_mask - def forward(self, x, t, y, attention_mask): + def forward(self, x, t, y, attention_mask=None, use_fp16=False): """ Forward pass of DiT. x: (B, T, C, H, W) tensor of spatial inputs (images or latent representations of images) @@ -326,7 +326,8 @@ class DiT(nn.Module): self.h = num_patches_height = H // self.patch_size self.w = num_patches_width = W // self.patch_size self.t = num_tubes_length = T // self.patch_size_t # 4 // 1 - attention_mask = self.make_mask(attention_mask, x.dtype) + if attention_mask is not None: + attention_mask = self.make_mask(attention_mask, x.dtype) pos_embed = get_2d_sincos_pos_embed(self.hidden_size, num_patches_height)