mask loss

.gitignore

@@ -27,4 +27,5 @@ sample_videos/
 720*
 1024*
 debug*
-private*
+private*
+caption*

opensora/models/diffusion/diffusion/gaussian_diffusion_t2v.py

@@ -694,7 +694,7 @@ class GaussianDiffusion_T:
                 img = out["sample"]
 
     def _vb_terms_bpd(
-            self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
+            self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None, mask=1.0, 
     ):
         """
         Get a term for the variational lower-bound.
@@ -704,6 +704,7 @@ class GaussianDiffusion_T:
                  - 'output': a shape [N] tensor of NLLs or KLs.
                  - 'pred_xstart': the x_0 predictions.
         """
+        # import ipdb;ipdb.set_trace()
         true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
             x_start=x_start, x_t=x_t, t=t
         )
@@ -713,13 +714,13 @@ class GaussianDiffusion_T:
         kl = normal_kl(
             true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
         )
-        kl = mean_flat(kl) / np.log(2.0)
+        kl = mean_flat(kl * mask) / np.log(2.0)
 
         decoder_nll = -discretized_gaussian_log_likelihood(
             x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
         )
         assert decoder_nll.shape == x_start.shape
-        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+        decoder_nll = mean_flat(decoder_nll * mask) / np.log(2.0)
 
         # At the first timestep return the decoder NLL,
         # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
@@ -740,12 +741,15 @@ class GaussianDiffusion_T:
         """
         if model_kwargs is None:
             model_kwargs = {}
+            mask = 1.0
+        else:
+            mask = model_kwargs['attention_mask'].unsqueeze(1)  # b t h w -> b 1 t h w
+
         if noise is None:
             noise = th.randn_like(x_start)
         x_t = self.q_sample(x_start, t, noise=noise)
 
         terms = {}
-
         if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
             terms["loss"] = self._vb_terms_bpd(
                 model=model,
@@ -786,7 +790,8 @@ class GaussianDiffusion_T:
                     x_start=x_start,
                     x_t=x_t,
                     t=t,
-                    clip_denoised=False,
+                    clip_denoised=False, 
+                    mask=mask
                 )["output"]
                 if self.loss_type == LossType.RESCALED_MSE:
                     # Divide by 1000 for equivalence with initial implementation.
@@ -801,7 +806,8 @@ class GaussianDiffusion_T:
                 ModelMeanType.EPSILON: noise,
             }[self.model_mean_type]
             assert model_output.shape == target.shape == x_start.shape
-            terms["mse"] = mean_flat((target - model_output) ** 2)
+            terms["mse"] = mean_flat(((target - model_output) ** 2) * mask)
+            # import ipdb;ipdb.set_trace()
             if "vb" in terms:
                 terms["loss"] = terms["mse"] + terms["vb"]
             else:

opensora/models/diffusion/latte/modeling_latte.py

@@ -256,25 +256,21 @@ class LatteT2V(ModelMixin, ConfigMixin):
         return pos_hw, pos_t
 
     def make_attn_mask(self, attention_mask, frame, dtype):
-        attention_mask = attention_mask.flatten(-2)  # *, h w -> *, hw
-        if attention_mask is not None and attention_mask.ndim == 2:
-            # assume that mask is expressed as:
-            #   (1 = keep,      0 = discard)
-            # convert mask into a bias that can be added to attention scores:
-            #   (keep = +0,     discard = -10000.0)
-            attention_mask = (1 - attention_mask.to(dtype)) * -10000.0
-            attention_mask = attention_mask.unsqueeze(1)
-            attention_mask = repeat(attention_mask, 'b 1 l -> (b f) 1 l', f=frame).contiguous()
-            attention_mask = attention_mask.to(self.dtype)
-        elif attention_mask is not None and attention_mask.ndim == 3:  # ndim == 3 means image joint
-            attention_mask = (1 - attention_mask.to(dtype)) * -10000.0
-            attention_mask_video = attention_mask[:, :1, ...]
-            attention_mask_video = repeat(attention_mask_video, 'b 1 l -> b (1 f) l',
-                                                  f=frame).contiguous()
-            attention_mask_image = attention_mask[:, 1:, ...]
-            attention_mask = torch.cat([attention_mask_video, attention_mask_image], dim=1)
-            attention_mask = rearrange(attention_mask, 'b n l -> (b n) l').contiguous().unsqueeze(1)
-            attention_mask = attention_mask.to(self.dtype)
+        attention_mask = rearrange(attention_mask, 'b t h w -> (b t) 1 (h w)') 
+        # assume that mask is expressed as:
+        #   (1 = keep,      0 = discard)
+        # convert mask into a bias that can be added to attention scores:
+        #   (keep = +0,     discard = -10000.0)
+        attention_mask = (1 - attention_mask.to(dtype)) * -10000.0
+        attention_mask = attention_mask.to(self.dtype)
+        return attention_mask
+
+    def vae_to_diff_mask(self, attention_mask, use_image_num):
+        dtype = attention_mask.dtype
+        # b, t+use_image_num, h, w, assume t as channel
+        # this version do not use 3d patch embedding
+        attention_mask = F.max_pool2d(attention_mask, kernel_size=(self.patch_size, self.patch_size), stride=(self.patch_size, self.patch_size))
+        attention_mask = attention_mask.bool().to(dtype)
         return attention_mask
 
     def forward(
@@ -343,16 +339,10 @@ class LatteT2V(ModelMixin, ConfigMixin):
         #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
         #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
         if attention_mask is None:
-            if use_image_num == 0:
-                attention_mask = torch.ones((input_batch_size, h//self.patch_size, w//self.patch_size), device=hidden_states.device, dtype=hidden_states.dtype)
-            else:
-                attention_mask = torch.ones((input_batch_size, 1+use_image_num, h//self.patch_size, w//self.patch_size), device=hidden_states.device, dtype=hidden_states.dtype)
-
+            attention_mask = torch.ones((input_batch_size, frame+use_image_num, h, w), device=hidden_states.device, dtype=hidden_states.dtype)
+        attention_mask = self.vae_to_diff_mask(attention_mask, use_image_num)
         dtype = attention_mask.dtype
-        if use_image_num == 0:
-            attention_mask_compress = F.max_pool2d(attention_mask.unsqueeze(1).float(), kernel_size=self.compress_kv_factor, stride=self.compress_kv_factor).squeeze(1)
-        else:
-            attention_mask_compress = F.max_pool2d(attention_mask.float(), kernel_size=self.compress_kv_factor, stride=self.compress_kv_factor)
+        attention_mask_compress = F.max_pool2d(attention_mask.float(), kernel_size=self.compress_kv_factor, stride=self.compress_kv_factor)
         attention_mask_compress = attention_mask_compress.to(dtype)
 
         attention_mask = self.make_attn_mask(attention_mask, frame, hidden_states.dtype)

opensora/train/train_t2v.py

@@ -161,7 +161,7 @@ def main(args):
         ae.vae.enable_tiling()
         ae.vae.tile_overlap_factor = args.tile_overlap_factor
         
-    kwargs = {'load_in_8bit': False, 'torch_dtype': weight_dtype}
+    kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True}
     text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval()
 
     ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
@@ -535,7 +535,7 @@ def main(args):
 
                     if args.enable_tracker:
                         validation_prompt = "The majestic beauty of a waterfall cascading down a cliff into a serene lake."
-                        logger.info(f"Running validation... \n"
+                        logger.info(f"Running validation...Using DDPM naive sampling...\n"
                                     f"Generating {args.num_validation_videos} videos with prompt: {validation_prompt}")
                         with torch.no_grad():
                             # create pipeline
@@ -615,6 +615,7 @@ if __name__ == "__main__":
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--model_max_length", type=int, default=300)
 
+    parser.add_argument('--enable_8bit_t5', action='store_true')
     parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
     parser.add_argument('--enable_tiling', action='store_true')
     parser.add_argument("--compress_kv", action="store_true")

opensora/utils/dataset_utils.py

@@ -79,21 +79,19 @@ class Collate:
         if self.use_image_num == 0:
             pad_batch_tubes, attention_mask = self.process(batch_tubes_vid, t_ds_stride, ds_stride, 
                                                       self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
-            # attention_mask: b t hw, use first frame
-            attention_mask = attention_mask[:, 0]  # b h w
+            # attention_mask: b t h w
             input_ids, cond_mask = input_ids_vid.squeeze(1), cond_mask_vid.squeeze(1)  # b 1 l -> b l
         else:
             pad_batch_tubes_vid, attention_mask_vid = self.process(batch_tubes_vid, t_ds_stride, ds_stride, 
                                                                    self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
-            # attention_mask_vid: b t hw, use first frame
-            attention_mask_vid = attention_mask_vid[:, :1]  # b 1 h w
+            # attention_mask_vid: b t h w
             pad_batch_tubes_img, attention_mask_img = self.process(batch_tubes_img, 1, ds_stride, 
                                                                    self.max_1hw, self.ae_stride_1hw, self.patch_size_1hw, extra_1=False)
             pad_batch_tubes_img = rearrange(pad_batch_tubes_img, '(b i) c 1 h w -> b c i h w', i=self.use_image_num)
             attention_mask_img = rearrange(attention_mask_img, '(b i) 1 h w -> b i h w', i=self.use_image_num)
             pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=2)  # concat at temporal, video first
-            # attention_mask_img: b num_img hw
-            attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1)  # b 1+num_img h w
+            # attention_mask_img: b num_img h w
+            attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1)  # b t+num_img h w
             input_ids = torch.cat([input_ids_vid, input_ids_img], dim=1)  # b 1+num_img hw
             cond_mask = torch.cat([cond_mask_vid, cond_mask_img], dim=1)  # b 1+num_img hw
         return pad_batch_tubes, attention_mask, input_ids, cond_mask
@@ -116,21 +114,28 @@ class Collate:
         pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0)
 
         # make attention_mask
-        max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
-        max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
-                           max_tube_size[1] // ae_stride_thw[1],
-                           max_tube_size[2] // ae_stride_thw[2]]
-        max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]),
-                                    max_latent_size[1] // patch_size_thw[1],
-                                    max_latent_size[2] // patch_size_thw[2]]
-        valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)),
-                                       int(math.ceil(i[2] / ds_stride)),
-                                       int(math.ceil(i[3] / ds_stride))] for i in batch_input_size]
-        attention_mask = [F.pad(torch.ones(i),
-                                (0, max_patchify_latent_size[2] - i[2],
-                                 0, max_patchify_latent_size[1] - i[1],
-                                 0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size]
-        attention_mask = torch.stack(attention_mask)
-        # attention_mask = attention_mask.flatten(-2)  # b t h w -> b t hw
-        attention_mask = attention_mask  # b t h w
+        first_channel_first_frame, first_channel_other_frame = pad_batch_tubes[:, :1, :1], pad_batch_tubes[:, :1, 1:]  # first channel to make attention_mask
+        attention_mask_first_frame = F.max_pool3d(first_channel_first_frame, kernel_size=(1, *ae_stride_thw[1:]), stride=(1, *ae_stride_thw[1:]))
+        if first_channel_other_frame.numel() != 0:
+            attention_mask_other_frame = F.max_pool3d(first_channel_other_frame, kernel_size=ae_stride_thw, stride=ae_stride_thw)
+            attention_mask = torch.cat([attention_mask_first_frame, attention_mask_other_frame], dim=2)
+        else:
+            attention_mask = attention_mask_first_frame
+        attention_mask = attention_mask[:, 0].bool().float()  # b t h w, do not channel
+        # max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
+        # max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
+        #                    max_tube_size[1] // ae_stride_thw[1],
+        #                    max_tube_size[2] // ae_stride_thw[2]]
+        # max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]),
+        #                             max_latent_size[1] // patch_size_thw[1],
+        #                             max_latent_size[2] // patch_size_thw[2]]
+        # valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)),
+        #                                int(math.ceil(i[2] / ds_stride)),
+        #                                int(math.ceil(i[3] / ds_stride))] for i in batch_input_size]
+        # attention_mask = [F.pad(torch.ones(i),
+        #                         (0, max_patchify_latent_size[2] - i[2],
+        #                          0, max_patchify_latent_size[1] - i[1],
+        #                          0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size]
+        # attention_mask = torch.stack(attention_mask)  # b t h w
+
         return pad_batch_tubes, attention_mask

scripts/text_condition/sample_video.sh

@@ -1,5 +1,5 @@
 CUDA_VISIBLE_DEVICES=6 python opensora/sample/sample_t2v.py \
-    --model_path /remote-home1/yeyang/sora4.8/Open-Sora-Plan/cache_dir/models--LanguageBind--Open-Sora-Plan-v1.0.0/snapshots/92a849c00ab99ae2c7cd7ff156f2ae02eb25c45d/65x512x512 \
+    --model_path /remote-home1/yeyang/dev/Open-Sora-Plan/debug_mask_loss/checkpoint-500/model \
     --version 65x512x512 \
     --image_size 512 \
     --cache_dir "/remote-home1/yeyang/sora4.8/Open-Sora-Plan/cache_dir" \