Delete free_lunch_utils.py

free_lunch_utils.py

@@ -1,340 +0,0 @@
-import torch
-import torch.fft as fft
-from diffusers.models.unets.unet_2d_condition import logger
-from diffusers.utils import is_torch_version
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-
-def isinstance_str(x: object, cls_name: str):
-    """
-    Checks whether x has any class *named* cls_name in its ancestry.
-    Doesn't require access to the class's implementation.
-    
-    Useful for patching!
-    """
-
-    for _cls in x.__class__.__mro__:
-        if _cls.__name__ == cls_name:
-            return True
-    
-    return False
-
-
-def Fourier_filter(x, threshold, scale):
-    dtype = x.dtype
-    x = x.type(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() 
-
-    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
-
-
-def register_upblock2d(model):
-    def up_forward(self):
-        def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
-            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:
-
-                    def create_custom_forward(module):
-                        def custom_forward(*inputs):
-                            return module(*inputs)
-
-                        return custom_forward
-
-                    if is_torch_version(">=", "1.11.0"):
-                        hidden_states = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                        )
-                    else:
-                        hidden_states = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(resnet), hidden_states, temb
-                        )
-                else:
-                    hidden_states = resnet(hidden_states, temb)
-
-            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, "UpBlock2D"):
-            upsample_block.forward = up_forward(upsample_block)
-
-
-def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
-    def up_forward(self):
-        def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
-            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 free upblock2d, hidden states shape: {hidden_states.shape}")
-                
-                # # --------------- 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)
-                # # ---------------------------------------------------------
-
-                # --------------- FreeU code -----------------------
-                # Only operate on the first two stages
-                if hidden_states.shape[1] == 1280:
-                    hidden_mean = hidden_states.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)
-                    
-                    hidden_states[:,:640] = hidden_states[:,:640] * ((self.b1 - 1 ) * hidden_mean + 1)
-                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
-                if hidden_states.shape[1] == 640:
-                    hidden_mean = hidden_states.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)
-                    
-                    hidden_states[:,:320] = hidden_states[:,:320] * ((self.b2 - 1 ) * hidden_mean + 1)
-                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
-                # ---------------------------------------------------------
-
-                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
-                if self.training and self.gradient_checkpointing:
-
-                    def create_custom_forward(module):
-                        def custom_forward(*inputs):
-                            return module(*inputs)
-
-                        return custom_forward
-
-                    if is_torch_version(">=", "1.11.0"):
-                        hidden_states = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                        )
-                    else:
-                        hidden_states = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(resnet), hidden_states, temb
-                        )
-                else:
-                    hidden_states = resnet(hidden_states, temb)
-
-            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, "UpBlock2D"):
-            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_upblock2d(model):
-    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,
-            encoder_attention_mask: Optional[torch.FloatTensor] = None,
-        ):
-            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)
-
-                if self.training and self.gradient_checkpointing:
-
-                    def create_custom_forward(module, return_dict=None):
-                        def custom_forward(*inputs):
-                            if return_dict is not None:
-                                return module(*inputs, return_dict=return_dict)
-                            else:
-                                return module(*inputs)
-
-                        return custom_forward
-
-                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(resnet),
-                        hidden_states,
-                        temb,
-                        **ckpt_kwargs,
-                    )
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(attn, return_dict=False),
-                        hidden_states,
-                        encoder_hidden_states,
-                        None,  # timestep
-                        None,  # class_labels
-                        cross_attention_kwargs,
-                        attention_mask,
-                        encoder_attention_mask,
-                        **ckpt_kwargs,
-                    )[0]
-                else:
-                    hidden_states = resnet(hidden_states, temb)
-                    hidden_states = attn(
-                        hidden_states,
-                        encoder_hidden_states=encoder_hidden_states,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                        attention_mask=attention_mask,
-                        encoder_attention_mask=encoder_attention_mask,
-                        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, "CrossAttnUpBlock2D"):
-            upsample_block.forward = up_forward(upsample_block)
-
-
-def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
-    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,
-            encoder_attention_mask: Optional[torch.FloatTensor] = None,
-        ):
-            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]
-
-                # --------------- FreeU code -----------------------
-                # Only operate on the first two stages
-                if hidden_states.shape[1] == 1280:
-                    hidden_mean = hidden_states.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)
-                    
-                    hidden_states[:,:640] = hidden_states[:,:640] * ((self.b1 - 1 ) * hidden_mean + 1)
-                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
-                if hidden_states.shape[1] == 640:
-                    hidden_mean = hidden_states.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)
-                    
-                    hidden_states[:,:320] = hidden_states[:,:320] * ((self.b2 - 1 ) * hidden_mean + 1)
-                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
-                # ---------------------------------------------------------
-
-                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
-                if self.training and self.gradient_checkpointing:
-
-                    def create_custom_forward(module, return_dict=None):
-                        def custom_forward(*inputs):
-                            if return_dict is not None:
-                                return module(*inputs, return_dict=return_dict)
-                            else:
-                                return module(*inputs)
-
-                        return custom_forward
-
-                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(resnet),
-                        hidden_states,
-                        temb,
-                        **ckpt_kwargs,
-                    )
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(attn, return_dict=False),
-                        hidden_states,
-                        encoder_hidden_states,
-                        None,  # timestep
-                        None,  # class_labels
-                        cross_attention_kwargs,
-                        attention_mask,
-                        encoder_attention_mask,
-                        **ckpt_kwargs,
-                    )[0]
-                else:
-                    hidden_states = resnet(hidden_states, temb)
-                    # hidden_states = attn(
-                    #     hidden_states,
-                    #     encoder_hidden_states=encoder_hidden_states,
-                    #     cross_attention_kwargs=cross_attention_kwargs,
-                    #     encoder_attention_mask=encoder_attention_mask,
-                    #     return_dict=False,
-                    # )[0]
-                    hidden_states = attn(
-                        hidden_states,
-                        encoder_hidden_states=encoder_hidden_states,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                    )[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, "CrossAttnUpBlock2D"):
-            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)