mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
salt
95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
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
|