scaled

modules_forge/packages/comfy/weight_adapter/oftv2.py

@@ -15,6 +15,7 @@ class OFTRotationUtil:
         eps: float = 6e-5,
         use_cayley_neumann: bool = True,
         num_cayley_neumann_terms: int = 5,
+        scaled_oft: bool = False,
     ):
         self.weight = weight
         self.block_size = block_size
@@ -22,6 +23,7 @@ class OFTRotationUtil:
         self.eps = eps
         self.use_cayley_neumann = use_cayley_neumann
         self.num_cayley_neumann_terms = num_cayley_neumann_terms
+        self.use_scaled_oft = scaled_oft
         self.rows, self.cols = torch.triu_indices(self.block_size, self.block_size, 1)
 
     def _get_triu_indices(self, device):
@@ -65,9 +67,11 @@ class OFTRotationUtil:
                     Q_squared = torch.bmm(Q_skew, Q_skew)
                     R.add_(Q_squared, alpha=2.0)
                     Q_power = Q_squared
-                    for _ in range(3, self.num_cayley_neumann_terms):
+                    for _ in range(3, self.num_cayley_neumann_terms - 1):
                         Q_power = torch.bmm(Q_power, Q_skew)
                         R.add_(Q_power, alpha=2.0)
+                    Q_power = torch.bmm(Q_power, Q_skew)
+                    R.add_(Q_power)
         else:
             id_mat = torch.eye(self.block_size, device=Q_skew.device).unsqueeze(0).expand_as(Q_skew)
             R = torch.linalg.solve(id_mat + Q_skew, id_mat - Q_skew, left=False)
@@ -79,7 +83,13 @@ class OFTRotationUtil:
             with torch.no_grad():
                 projected_weight = self._project_batch()
                 weight.copy_(projected_weight)
-        return self._cayley_batch(weight)
+        if self.use_scaled_oft:
+            # Apply the scaled_oft factor: weight / sqrt(n_elements)
+            n_elements = weight.shape[-1]
+            effective_weight = weight / (n_elements**0.5)
+        else:
+            effective_weight = weight
+        return self._cayley_batch(effective_weight)
 
 
 class OFTv2Adapter(WeightAdapterBase):
@@ -101,10 +111,15 @@ class OFTv2Adapter(WeightAdapterBase):
         if loaded_keys is None:
             loaded_keys = set()
         oft_r_weight_name = f"{x}.oft_R.weight"
+        scaled_oft_flag_name = f"{x}.oft_R.scaled_oft"
         if oft_r_weight_name in lora:
             oft_r_weight = lora[oft_r_weight_name]
             loaded_keys.add(oft_r_weight_name)
-            weights = (oft_r_weight, alpha, dora_scale)
+            is_scaled = False
+            if scaled_oft_flag_name in lora:
+                is_scaled = True
+                loaded_keys.add(scaled_oft_flag_name)
+            weights = (oft_r_weight, alpha, dora_scale, is_scaled)
             return cls(loaded_keys, weights)
         return None
 
@@ -122,7 +137,7 @@ class OFTv2Adapter(WeightAdapterBase):
         if strength == 0.0:
             return weight
 
-        oft_r_weight_orig, alpha, dora_scale = self.weights
+        oft_r_weight_orig, alpha, dora_scale, is_scaled = self.weights
 
         try:
             oft_r_weight_processed = oft_r_weight_orig.to(weight.device, dtype=intermediate_dtype)
@@ -150,7 +165,7 @@ class OFTv2Adapter(WeightAdapterBase):
                 return weight
 
             # Pass the unscaled weight to the utility to get the full rotation matrix
-            util = OFTRotationUtil(oft_r_weight_processed, block_size)
+            util = OFTRotationUtil(oft_r_weight_processed, block_size, scaled_oft=is_scaled)
             orth_rotate = util.get_rotation_matrix()
 
             # For Linear layers,  rotates the input (x @ R), equivalent to rotating weights by R.T (W @ R.T).