From 5b5cabe33d46a5b1828e59b8076514d6d427b0aa Mon Sep 17 00:00:00 2001 From: tin2tin Date: Thu, 28 Sep 2023 07:14:55 +0200 Subject: [PATCH] Add FreeU for Zeroscope --- free_lunch_utils.py | 310 ++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 283 insertions(+), 27 deletions(-) diff --git a/free_lunch_utils.py b/free_lunch_utils.py index 005b4f7..c68c947 100644 --- a/free_lunch_utils.py +++ b/free_lunch_utils.py @@ -1,10 +1,10 @@ -# from https://github.com/lyn-rgb/FreeU_Diffusers +from typing import Any, Dict, Optional, Tuple import torch import torch.fft as fft -from diffusers.models.unet_2d_condition import logger from diffusers.utils import is_torch_version -from typing import Any, Dict, List, Optional, Tuple, Union +from diffusers.models.unet_2d_condition import logger as logger2d +from diffusers.models.unet_3d_condition import logger as logger3d def isinstance_str(x: object, cls_name: str): @@ -22,36 +22,57 @@ def isinstance_str(x: object, cls_name: str): return False -def Fourier_filter(x, threshold, scale): - dtype = x.dtype - x = x.type(torch.float32) +def Fourier_filter(x_in, threshold, scale): + """ + Updated Fourier filter based on: + https://github.com/huggingface/diffusers/pull/5164#issuecomment-1732638706 + """ + + x = x_in + B, C, H, W = x.shape + + # Non-power of 2 images must be float32 + if (W & (W - 1)) != 0 or (H & (H - 1)) != 0: + x = x.to(dtype=torch.float32) + # FFT 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() + mask = torch.ones((B, C, H, W), device=x.device) - crow, ccol = H // 2, W //2 - mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = scale + crow, ccol = H // 2, W // 2 + mask[..., crow - threshold : crow + threshold, ccol - threshold : ccol + threshold] = scale x_freq = x_freq * mask # IFFT x_freq = fft.ifftshift(x_freq, dim=(-2, -1)) x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real - - x_filtered = x_filtered.type(dtype) - return x_filtered + + return x_filtered.to(dtype=x_in.dtype) def register_upblock2d(model): + """ + Register UpBlock2D for UNet2DCondition. + """ + def up_forward(self): - def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None): + def forward( + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None, + scale: float = 1.0 + ): + logger2d.debug(f"in upblock2d, hidden states shape: {hidden_states.shape}") + for resnet in self.resnets: # pop res hidden states res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] - #print(f"in upblock2d, hidden states shape: {hidden_states.shape}") + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) if self.training and self.gradient_checkpointing: @@ -71,11 +92,11 @@ def register_upblock2d(model): create_custom_forward(resnet), hidden_states, temb ) else: - hidden_states = resnet(hidden_states, temb) + hidden_states = resnet(hidden_states, temb, scale=scale) if self.upsamplers is not None: for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) return hidden_states @@ -87,14 +108,25 @@ def register_upblock2d(model): def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): + """ + Register UpBlock2D with FreeU for UNet2DCondition. + """ + def up_forward(self): - def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + def forward( + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None, + scale: float = 1.0 + ): + logger2d.debug(f"in free upblock2d, hidden states shape: {hidden_states.shape}") + for resnet in self.resnets: # pop res hidden states - #print(f"in free upblock2d, hidden states shape: {hidden_states.shape}") res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] - + # --------------- FreeU code ----------------------- # Only operate on the first two stages if hidden_states.shape[1] == 1280: @@ -144,6 +176,10 @@ def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): def register_crossattn_upblock2d(model): + """ + Register CrossAttn UpBlock2D for UNet2DCondition. + """ + def up_forward(self): def forward( hidden_states: torch.FloatTensor, @@ -155,9 +191,12 @@ def register_crossattn_upblock2d(model): attention_mask: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, ): + logger2d.debug(f"in crossatten upblock2d, hidden states shape: {hidden_states.shape}") + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + for resnet, attn in zip(self.resnets, self.attentions): # pop res hidden states - #print(f"in crossatten upblock2d, hidden states shape: {hidden_states.shape}") res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) @@ -192,7 +231,7 @@ def register_crossattn_upblock2d(model): **ckpt_kwargs, )[0] else: - hidden_states = resnet(hidden_states, temb) + hidden_states = resnet(hidden_states, temb, scale=lora_scale) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, @@ -204,7 +243,7 @@ def register_crossattn_upblock2d(model): if self.upsamplers is not None: for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) return hidden_states @@ -216,6 +255,10 @@ def register_crossattn_upblock2d(model): def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): + """ + Register CrossAttn UpBlock2D with FreeU for UNet2DCondition. + """ + def up_forward(self): def forward( hidden_states: torch.FloatTensor, @@ -227,9 +270,12 @@ def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): attention_mask: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, ): + logger2d.debug(f"in free crossatten upblock2d, hidden states shape: {hidden_states.shape}") + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + for resnet, attn in zip(self.resnets, self.attentions): # pop res hidden states - #print(f"in free crossatten upblock2d, hidden states shape: {hidden_states.shape}") res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] @@ -275,7 +321,7 @@ def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): **ckpt_kwargs, )[0] else: - hidden_states = resnet(hidden_states, temb) + hidden_states = resnet(hidden_states, temb, scale=lora_scale) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, @@ -287,7 +333,7 @@ def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): if self.upsamplers is not None: for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) return hidden_states @@ -299,4 +345,214 @@ def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): setattr(upsample_block, 'b1', b1) setattr(upsample_block, 'b2', b2) setattr(upsample_block, 's1', s1) - setattr(upsample_block, 's2', s2) \ No newline at end of file + setattr(upsample_block, 's2', s2) + + +def register_upblock3d(model): + """ + Register UpBlock3D for UNet3DCondition. + """ + + def up_forward(self): + def forward( + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None, + num_frames=1 + ): + + logger3d.debug(f"in upblock3d, hidden states shape: {hidden_states.shape}") + + for resnet, temp_conv in zip(self.resnets, self.temp_convs): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + return forward + + for i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "UpBlock3D"): + upsample_block.forward = up_forward(upsample_block) + + +def register_free_upblock3d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): + """ + Register UpBlock3D with FreeU for UNet3DCondition. + """ + + def up_forward(self): + def forward( + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None, + num_frames=1 + ): + + logger3d.debug(f"in free upblock3d, hidden states shape: {hidden_states.shape}") + + for resnet, temp_conv in zip(self.resnets, self.temp_convs): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # --------------- FreeU code ----------------------- + # 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) + 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) + # --------------------------------------------------------- + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + return forward + + for i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "UpBlock3D"): + upsample_block.forward = up_forward(upsample_block) + setattr(upsample_block, 'b1', b1) + setattr(upsample_block, 'b2', b2) + setattr(upsample_block, 's1', s1) + setattr(upsample_block, 's2', s2) + + +def register_crossattn_upblock3d(model): + """ + Register CrossAttn UpBlock3D for UNet3DCondition. + """ + + def up_forward(self): + def forward( + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1 + ): + logger3d.debug(f"in crossatten upblock3d, hidden states shape: {hidden_states.shape}") + + for resnet, temp_conv, attn, temp_attn in zip( + self.resnets, self.temp_convs, self.attentions, self.temp_attentions + ): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = temp_attn( + hidden_states, num_frames=num_frames, cross_attention_kwargs=cross_attention_kwargs, return_dict=False + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + return forward + + for i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "CrossAttnUpBlock3D"): + upsample_block.forward = up_forward(upsample_block) + + +def register_free_crossattn_upblock3d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2): + """ + Register CrossAttn UpBlock3D with FreeU for UNet3DCondition. + """ + + def up_forward(self): + def forward( + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1 + ): + logger3d.debug(f"in free crossatten upblock3d, hidden states shape: {hidden_states.shape}") + + for resnet, temp_conv, attn, temp_attn in zip( + self.resnets, self.temp_convs, self.attentions, self.temp_attentions + ): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # --------------- FreeU code ----------------------- + # 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) + 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) + # --------------------------------------------------------- + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = temp_attn( + hidden_states, num_frames=num_frames, cross_attention_kwargs=cross_attention_kwargs, return_dict=False + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + return forward + + for i, upsample_block in enumerate(model.unet.up_blocks): + if isinstance_str(upsample_block, "CrossAttnUpBlock3D"): + upsample_block.forward = up_forward(upsample_block) + setattr(upsample_block, 'b1', b1) + setattr(upsample_block, 'b2', b2) + setattr(upsample_block, 's1', s1) + setattr(upsample_block, 's2', s2)