Update model.py

model.py

@@ -133,29 +133,35 @@ class rotary(nn.Module):
     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.dims = dims
+        self.use_pbias = False
+
         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.use_pbias = False    
+        self.dims = dims
         self.max_ctx = max_ctx
         self.variable_radius = variable_radius
         
-        self.inv_freq = nn.Parameter(1.0 / (theta ** (torch.arange(0, dims, 2, device=device, dtype=dtype) / dims)),
+        self.inv_freq = nn.Parameter(
+                1.0 / (10000 ** (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.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.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)
     
         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:
@@ -194,27 +200,38 @@ class rotary(nn.Module):
             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()
-            f0_theta = (f0_mean**2) * self.pitch_scale
-            #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)) 
+            f0_mean = torch.clamp(f0_mean, min=80.0, max=600.0)
+            perceptual_factor = torch.log(1 + f0_mean / 700.0) / torch.log(torch.tensor(1 + 300.0 / 700.0))
+            f0_theta = self.min_theta + perceptual_factor * (self.max_theta - self.min_theta)
+            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:
-            if f0 is not None:
-                f0 = f0[0]
-                seq_len = x
-                f0 = torch.tensor(f0, device=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 f0 is not None:
+#                 f0 = f0[0]
+#                 seq_len = x
+#                 f0 = self.align_f0_to_tokens(f0, freqs.shape[-1])
+#                 radius = f0
+          
+#                 freqs = torch.polar(radius, freqs)
+#             else:
+          
+#                 freqs = torch.polar(torch.ones_like(freqs), freqs)
+#         freqs = freqs.unsqueeze(0)
+
+            radius = F.softplus(self.radius)
+            freqs = torch.polar(radius.unsqueeze(0).expand_as(freqs), 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}"
@@ -222,12 +239,13 @@ 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) > 0.0:
+                    elif abs(self._prev_f0_theta - f0_theta) > 200.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,11 +258,13 @@ 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)