Update model.py

model.py

@@ -130,38 +130,32 @@ def sinusoids(length, channels, max_timescale=10000):
 
 class rotary(nn.Module):
     _seen = set()  
-    def __init__(self, dims, max_ctx=1500, theta=10000, learned_freq=True, variable_radius=False,
+    def __init__(self, dims, max_ctx=1500, theta=10000, learned_freq=False, variable_radius=False,
                  learned_radius=False, learned_theta=False, learned_pitch=False, debug: List[str] = []):
         super().__init__()
-        self.use_pbias = False
-
+        
+        self.dims = dims
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        dtype = torch.float32
         self.device = device
-        dtype = torch.float32 
         self.dtype = dtype
         self.debug = debug
         self._counter = 0
-        self.dims = dims
+
+        self.use_pbias = False    
         self.max_ctx = max_ctx
         self.variable_radius = variable_radius
         
-        self.inv_freq = nn.Parameter(
-                1.0 / (19000 ** (torch.arange(0, dims, 2, device=device, dtype=dtype) / dims)),
+        self.inv_freq = nn.Parameter(1.0 / (theta ** (torch.arange(0, dims, 2, device=device, dtype=dtype) / dims)),
                 requires_grad=learned_freq)
-        self.theta = nn.Parameter(
-            torch.tensor(float(theta)), requires_grad=learned_theta)
-        self.min_theta = nn.Parameter(
-            torch.tensor(600.0), requires_grad=learned_theta)
-        self.max_theta = nn.Parameter(
-            torch.tensor(2400.0), requires_grad=learned_theta)
-        
-        self.pitch_scale = nn.Parameter(torch.tensor(1.0), 
-                                        requires_grad=learned_pitch)
+
+        self.theta = nn.Parameter(torch.tensor(float(theta)), 
+                requires_grad=learned_theta)
+
+        self.pitch_scale = nn.Parameter(torch.tensor(1.0), requires_grad=learned_pitch)
     
         if variable_radius:
-            self.radius = nn.Parameter(
-                torch.ones(dims // 2),
-                requires_grad=learned_radius)
+            self.radius = nn.Parameter(torch.ones(dims // 2), requires_grad=learned_radius)
 
     def get_pitch_bias(self, f0):
         if f0 is None:
@@ -189,31 +183,38 @@ class rotary(nn.Module):
         rotary.get_sim = get_sim
         rotary.fwd_sim = fwd_sim
 
-    def forward(self, x = None, f0=None) -> Tensor:
+    def align_f0_to_tokens(self, f0, token_length):
+        ratio = len(f0) / token_length
+        indices = [int(i * ratio) for i in range(token_length)]
+        indices = [min(i, len(f0) - 1) for i in indices]
+        return f0[indices]
+
+    def forward(self, x=None, f0=None, stage=None) -> Tensor:
         if isinstance(x, int):
             t = torch.arange(x, device=self.device).float()
         else:
             t = x.float().to(self.inv_freq.device)
-
         if f0 is not None:
-                
             f0_mean = f0.mean()
-            perceptual_factor = torch.log(1 + f0_mean / 700.0) / torch.log(torch.tensor(1 + 300.0 / 700.0))
-            min_theta, max_theta = 800.0, 10000.0
-            f0_theta = self.theta + perceptual_factor * (max_theta - min_theta)
-            inv_freq = 1.0 / (f0_theta ** (torch.arange(0, self.dims, 2, device=self.device) / self.dims))
+            f0_theta = f0_mean * (f0_mean / self.theta) * self.theta * self.pitch_scale
+            inv_freq = 1.0 / (f0_theta ** (torch.arange(0, self.dims, 2, device=self.device) / self.dims)) 
         else:
             inv_freq = self.inv_freq
         freqs = torch.einsum('i,j->ij', t, inv_freq)
-
         freqs = freqs.float()
         if self.variable_radius:
-            radius = F.softplus(self.radius)
-            freqs = torch.polar(radius.unsqueeze(0).expand_as(freqs), freqs)
-        else:
-            freqs = torch.polar(torch.ones_like(freqs), freqs)
+
+            if f0 is not None:
+                f0 = f0[0]
+                seq_len = x
+                f0 = torch.tensor(f0, device=x.device if isinstance(x, torch.Tensor) else device)
+                f0 = self.align_f0_to_tokens(f0, freqs.shape[-1])
+                radius = 1.0 / (f0 + 1)
+                freqs = torch.polar(radius, freqs)
+            else:
+                freqs = torch.polar(torch.ones_like(freqs), freqs)
         freqs = freqs.unsqueeze(0)
-            
+
         if "rotary" in self.debug:
             if f0 is not None:
                 key = f"{self._counter}_{f0_theta:.2f}"
@@ -221,13 +222,12 @@ class rotary(nn.Module):
                     if not hasattr(self, '_prev_f0_theta'):
                         self._prev_f0_theta = f0_theta
                         print(f"Step {self._counter}: Using raw F0 as theta: {f0_theta:.2f} Hz")
-                    elif abs(self._prev_f0_theta - f0_theta) > 200.0:
+                    elif abs(self._prev_f0_theta - f0_theta) > 0.0:
                         print(f"Step {self._counter}: Using raw F0 as theta: {f0_theta:.2f} Hz")
                         self._prev_f0_theta = f0_theta
                     rotary._seen.add(key)
             self._counter += 1
-            
-        return freqs      
+        return freqs
 
     @staticmethod
     def apply_rotary(x, freqs):
@@ -240,13 +240,11 @@ class rotary(nn.Module):
             x1 = x1 * freqs
             x1 = torch.view_as_real(x1).flatten(-2)
             return torch.cat([x1.type_as(x), x2], dim=-1)
-
         else:
             x1 = x[..., :freqs.shape[-1]*2]
             x2 = x[..., freqs.shape[-1]*2:]
             
             if x.ndim == 2:  
-  
                 x1 = x1.unsqueeze(0)
                 x1 = x1.float().reshape(*x1.shape[:-1], -1, 2).contiguous()
                 x1 = torch.view_as_complex(x1)
@@ -260,7 +258,7 @@ class rotary(nn.Module):
                 x1 = x1 * freqs
                 x1 = torch.view_as_real(x1).flatten(-2)
                 return torch.cat([x1.type_as(x), x2], dim=-1)
-
+                
 class SliceAttention(nn.Module):
     def __init__(self, dims, heads, dropout=0.0):
         super().__init__()