From 06ffc8eccfba0bb34c8827b8db41c398b941607e Mon Sep 17 00:00:00 2001
From: flymin <905370712@qq.com>
Date: Tue, 16 Jul 2024 11:01:20 +0800
Subject: [PATCH 1/3] fix bug at mha in blocks.py

---
 opensora/models/layers/blocks.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/opensora/models/layers/blocks.py b/opensora/models/layers/blocks.py
index 40e6abbc..e28ef343 100644
--- a/opensora/models/layers/blocks.py
+++ b/opensora/models/layers/blocks.py
@@ -467,6 +467,11 @@ def forward(self, x, cond, mask=None):
         # query/value: img tokens; key: condition; mask: if padding tokens
         B, N, C = x.shape
 
+        if mask is None:
+            Bc, Nc, _ = cond.shape
+            assert Bc == B
+            mask = [Nc] * B
+
         q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
         kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
         k, v = kv.unbind(2)
@@ -504,6 +509,11 @@ def forward(self, x, cond, mask=None):
         B, SUB_N, C = x.shape  # [B, TS/p, C]
         N = SUB_N * sp_size
 
+        if mask is None:
+            Bc, Nc, _ = cond.shape
+            assert Bc == B
+            mask = [Nc] * B
+
         # shape:
         # q, k, v: [B, SUB_N, NUM_HEADS, HEAD_DIM]
         q = self.q_linear(x).view(B, -1, self.num_heads, self.head_dim)

From 8af7f69ea22c18f8b59f3b290d5846f0524632d1 Mon Sep 17 00:00:00 2001
From: flymin <905370712@qq.com>
Date: Tue, 16 Jul 2024 11:02:03 +0800
Subject: [PATCH 2/3] fix bug in MaskGenerator

---
 opensora/utils/train_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/opensora/utils/train_utils.py b/opensora/utils/train_utils.py
index 95d0011b..21391d24 100644
--- a/opensora/utils/train_utils.py
+++ b/opensora/utils/train_utils.py
@@ -153,9 +153,9 @@ def get_mask(self, x):
         elif mask_name == "random":
             mask_ratio = random.uniform(0.1, 0.9)
             mask = torch.rand(num_frames, device=x.device) > mask_ratio
-            # if mask is all False, set the last frame to True
-            if not mask.any():
-                mask[-1] = 1
+        # if mask is all False, set the last frame to True
+        if not mask.any():
+            mask[-1] = 1
 
         return mask
 

From bbdb698ab40b93d874e1fa29189e8bf497504830 Mon Sep 17 00:00:00 2001
From: flymin <905370712@qq.com>
Date: Tue, 16 Jul 2024 11:03:37 +0800
Subject: [PATCH 3/3] align logging style in ckpt_utils.py

---
 opensora/utils/ckpt_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/opensora/utils/ckpt_utils.py b/opensora/utils/ckpt_utils.py
index d730981c..f8b0e3f3 100644
--- a/opensora/utils/ckpt_utils.py
+++ b/opensora/utils/ckpt_utils.py
@@ -197,8 +197,8 @@ def load_checkpoint(model, ckpt_path, save_as_pt=False, model_name="model", stri
         from safetensors.torch import load_file
         state_dict = load_file(ckpt_path)
         missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
-        print(f"Missing keys: {missing_keys}")
-        print(f"Unexpected keys: {unexpected_keys}")
+        get_logger().info("Missing keys: %s", missing_keys)
+        get_logger().info("Unexpected keys: %s", unexpected_keys)
     elif os.path.isdir(ckpt_path):
         load_from_sharded_state_dict(model, ckpt_path, model_name, strict=strict)
         get_logger().info("Model checkpoint loaded from %s", ckpt_path)