Update model.py

model.py

@@ -196,24 +196,30 @@ class rotary(nn.Module):
             t = x.float().to(self.inv_freq.device)
         if f0 is not None:
             f0_mean = f0.mean()
-            f0_theta = f0_mean * (f0_mean / self.theta) * self.theta * self.pitch_scale
+            f0_theta = (f0_mean**2) * 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:
-
             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 = torch.tensor(f0, device=x.device if isinstance(x, torch.Tensor) else self.device)
                 f0 = self.align_f0_to_tokens(f0, freqs.shape[-1])
-                radius = 1.0 / (f0 + 1)
+                max_f0 = torch.max(f0)
+ 
+                if max_f0 > 0:
+                    radius = f0 / max_f0
+                else:
+                    radius = torch.ones_like(f0)
+                    
                 freqs = torch.polar(radius, freqs)
             else:
                 freqs = torch.polar(torch.ones_like(freqs), freqs)
-        freqs = freqs.unsqueeze(0)
+            freqs = freqs.unsqueeze(0)
+            print(f"radius, {radius}")
 
         if "rotary" in self.debug:
             if f0 is not None:
@@ -258,7 +264,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__()