From 5e3a3c6f78f06f8ea05fe8a230deb4d330175446 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sat, 27 Apr 2024 14:53:23 +0800 Subject: [PATCH] mask loss --- .gitignore | 3 +- .../diffusion/gaussian_diffusion_t2v.py | 18 ++++--- .../models/diffusion/latte/modeling_latte.py | 46 +++++++---------- opensora/train/train_t2v.py | 5 +- opensora/utils/dataset_utils.py | 51 ++++++++++--------- scripts/text_condition/sample_video.sh | 2 +- 6 files changed, 64 insertions(+), 61 deletions(-) diff --git a/.gitignore b/.gitignore index 31d1e45b7..95d0f4204 100644 --- a/.gitignore +++ b/.gitignore @@ -27,4 +27,5 @@ sample_videos/ 720* 1024* debug* -private* \ No newline at end of file +private* +caption* \ No newline at end of file diff --git a/opensora/models/diffusion/diffusion/gaussian_diffusion_t2v.py b/opensora/models/diffusion/diffusion/gaussian_diffusion_t2v.py index d29da269b..b99b0d097 100644 --- a/opensora/models/diffusion/diffusion/gaussian_diffusion_t2v.py +++ b/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: diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index 263cc1a50..340acb6dc 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/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) diff --git a/opensora/train/train_t2v.py b/opensora/train/train_t2v.py index a74c45bc4..b18828f6b 100644 --- a/opensora/train/train_t2v.py +++ b/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") diff --git a/opensora/utils/dataset_utils.py b/opensora/utils/dataset_utils.py index f86821430..8ab656446 100644 --- a/opensora/utils/dataset_utils.py +++ b/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 diff --git a/scripts/text_condition/sample_video.sh b/scripts/text_condition/sample_video.sh index 9a39cf1ce..fb6b7bd10 100644 --- a/scripts/text_condition/sample_video.sh +++ b/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" \