nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fixed dit forward

Browse Source 
Former-commit-id: 4cd537fab12231f7471107762b02b36217ca42d8
This commit is contained in:
LinB203
2024-03-07 16:41:31 +08:00
parent a03025a8f3
commit 11996dc9b4
2 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
opensora/models/ae/videobase/vqvae/videogpt/rec_video.py
+2 -3
View File
@@ -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)
opensora/models/diffusion/dit/dit.py
+3 -2
View File
@@ -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)