From 75938cfaeb1fe0d95e073848af579a04fd20f547 Mon Sep 17 00:00:00 2001
From: jpthu17 <1003901354@qq.com>
Date: Sat, 9 Mar 2024 22:04:15 +0800
Subject: [PATCH] [fix] 2d RoPE init

Former-commit-id: ff5931bae6efa94675b9ef6bc25f1b76217d128e
---
 opensora/models/diffusion/dit/dit.py     | 8 +++++---
 opensora/models/diffusion/latte/latte.py | 8 +++++---
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py
index b6f3df4f5..e514a213a 100644
--- a/opensora/models/diffusion/dit/dit.py
+++ b/opensora/models/diffusion/dit/dit.py
@@ -49,7 +49,7 @@ class Attention(nn.Module):
             attention_pe_mode=None,
             hw: Union[int, Tuple[int, int]] = 16,  # (h, w)
             pt_hw: Union[int, Tuple[int, int]] = 16,  # (h, w)
-            intp_vfreq: bool = False,  # vision position interpolation
+            intp_vfreq: bool = True,  # vision position interpolation
     ) -> None:
         super().__init__()
         assert dim % num_heads == 0, 'dim should be divisible by num_heads'
@@ -280,8 +280,8 @@ class DiT(nn.Module):
         extras=1,
         attention_mode='math',
         attention_pe_mode=None,
-        pt_input_size: Union[int, Tuple[int, int]] = 16,  # (h, w)
-        intp_vfreq: bool = False,  # vision position interpolation
+        pt_input_size: Union[int, Tuple[int, int]] = None,  # (h, w)
+        intp_vfreq: bool = True,  # vision position interpolation
     ):
         super().__init__()
         self.gradient_checkpointing = False
@@ -305,6 +305,8 @@ class DiT(nn.Module):
         self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
         self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False)
 
+        if pt_input_size is None:
+            pt_input_size = input_size
         self.blocks = nn.ModuleList([
             DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode,
                      attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size,
diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py
index 42fdb5922..e89a87e57 100644
--- a/opensora/models/diffusion/latte/latte.py
+++ b/opensora/models/diffusion/latte/latte.py
@@ -56,7 +56,7 @@ class Attention(nn.Module):
                  attention_pe_mode=None,
                  hw: Union[int, Tuple[int, int]] = 16,  # (h, w)
                  pt_hw: Union[int, Tuple[int, int]] = 16,  # (h, w)
-                 intp_vfreq: bool = False,  # vision position interpolation
+                 intp_vfreq: bool = True,  # vision position interpolation
                  ):
         super().__init__()
         assert dim % num_heads == 0, 'dim should be divisible by num_heads'
@@ -284,8 +284,8 @@ class Latte(nn.Module):
         extras=1,
         attention_mode='math',
         attention_pe_mode=None,
-        pt_input_size: Union[int, Tuple[int, int]] = 16,  # (h, w)
-        intp_vfreq: bool = False,  # vision position interpolation
+        pt_input_size: Union[int, Tuple[int, int]] = None,  # (h, w)
+        intp_vfreq: bool = True,  # vision position interpolation
     ):
         super().__init__()
         self.learn_sigma = learn_sigma
@@ -316,6 +316,8 @@ class Latte(nn.Module):
         self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False)
         self.hidden_size = hidden_size
 
+        if pt_input_size is None:
+            pt_input_size = input_size
         self.blocks = nn.ModuleList([
             TransformerBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode,
                              attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size,