Update free_lunch_utils.py

free_lunch_utils.py

@@ -97,10 +97,10 @@ def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
  # Only operate on the first two stages
  if hidden_states.shape[1] == 1280:
  hidden_states[:,:640] = hidden_states[:,:640] * self.b1
- res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
+ # # res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
  if hidden_states.shape[1] == 640:
  hidden_states[:,:320] = hidden_states[:,:320] * self.b2
- res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
+ # res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
  # ---------------------------------------------------------
 
  hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
@@ -235,10 +235,10 @@ def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
  # Only operate on the first two stages
  if hidden_states.shape[1] == 1280:
  hidden_states[:,:640] = hidden_states[:,:640] * self.b1
- res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
+ # res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
  if hidden_states.shape[1] == 640:
  hidden_states[:,:320] = hidden_states[:,:320] * self.b2
- res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
+ # res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
  # ---------------------------------------------------------
 
  hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)