add rerender lora

animation/rerender-lora/src/freeu.py

@@ -0,0 +1,94 @@
+import torch
+import torch.fft as fft
+
+
+def Fourier_filter(x, threshold, scale):
+
+    x_freq = fft.fftn(x, dim=(-2, -1))
+    x_freq = fft.fftshift(x_freq, dim=(-2, -1))
+
+    B, C, H, W = x_freq.shape
+    mask = torch.ones((B, C, H, W)).cuda()
+
+    crow, ccol = H // 2, W // 2
+    mask[..., crow - threshold:crow + threshold,
+         ccol - threshold:ccol + threshold] = scale
+    x_freq = x_freq * mask
+
+    x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
+
+    x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
+
+    return x_filtered
+
+from deps.ControlNet.ldm.modules.diffusionmodules.util import \
+    timestep_embedding  # noqa：E501
+
+
+# backbone_scale1=1.1, backbone_scale2=1.2, skip_scale1=1.0, skip_scale2=0.2
+def freeu_forward(self,
+                  backbone_scale1=1.,
+                  backbone_scale2=1.,
+                  skip_scale1=1.,
+                  skip_scale2=1.):
+
+    def forward(x,
+                timesteps=None,
+                context=None,
+                control=None,
+                only_mid_control=False,
+                **kwargs):
+        hs = []
+        with torch.no_grad():
+            t_emb = timestep_embedding(timesteps,
+                                       self.model_channels,
+                                       repeat_only=False)
+            emb = self.time_embed(t_emb)
+            h = x.type(self.dtype)
+            for module in self.input_blocks:
+                h = module(h, emb, context)
+                hs.append(h)
+            h = self.middle_block(h, emb, context)
+
+        if control is not None:
+            h += control.pop()
+        '''
+        for i, module in enumerate(self.output_blocks):
+            if only_mid_control or control is None:
+                h = torch.cat([h, hs.pop()], dim=1)
+            else:
+                h = torch.cat([h, hs.pop() + control.pop()], dim=1)
+            h = module(h, emb, context)
+        '''
+        for i, module in enumerate(self.output_blocks):
+            hs_ = hs.pop()
+
+            if h.shape[1] == 1280:
+                hidden_mean = h.mean(1).unsqueeze(1)
+                B = hidden_mean.shape[0]
+                hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True) 
+                hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True)
+                hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3)
+                h[:, :640] = h[:, :640] * ((backbone_scale1 - 1) * hidden_mean + 1)        
+                # h[:, :640] = h[:, :640] * backbone_scale1
+                hs_ = Fourier_filter(hs_, threshold=1, scale=skip_scale1)
+            if h.shape[1] == 640:
+                hidden_mean = h.mean(1).unsqueeze(1)
+                B = hidden_mean.shape[0]
+                hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True) 
+                hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True)
+                hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3) 
+                h[:, :320] = h[:, :320] * ((backbone_scale2 - 1) * hidden_mean + 1)
+                # h[:, :320] = h[:, :320] * backbone_scale2
+                hs_ = Fourier_filter(hs_, threshold=1, scale=skip_scale2)
+
+            if only_mid_control or control is None:
+                h = torch.cat([h, hs_], dim=1)
+            else:
+                h = torch.cat([h, hs_ + control.pop()], dim=1)
+            h = module(h, emb, context)
+
+        h = h.type(x.dtype)
+        return self.out(h)
+
+    return forward