add freeU model

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14307648


* support sd21, sdxl

modelscope/metainfo.py

@@ -291,6 +291,7 @@ class Pipelines(object):
     image_denoise = 'nafnet-image-denoise'
     image_deblur = 'nafnet-image-deblur'
     image_editing = 'masactrl-image-editing'
+    freeu_stable_diffusion_text2image = 'freeu-stable-diffusion-text2image'
     person_image_cartoon = 'unet-person-image-cartoon'
     ocr_detection = 'resnet18-ocr-detection'
     table_recognition = 'dla34-table-recognition'

modelscope/models/multi_modal/freeu/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from modelscope.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+    from .free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
+else:
+    _import_structure = {
+        'free_lunch_utils':
+        ['register_free_upblock2d', 'register_free_crossattn_upblock2d']
+    }
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/models/multi_modal/freeu/free_lunch_utils.py

@@ -0,0 +1,331 @@
+# ------------------------------------------------------------------------
+# Modified from https://github.com/ChenyangSi/FreeU/blob/main/demo/free_lunch_utils.py
+# Copyright (c) 2023 TencentARC. All Rights Reserved.
+# ------------------------------------------------------------------------
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.fft as fft
+from diffusers.utils import is_torch_version
+
+
+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]
+                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]
+
+                # --------------- 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)
+
+                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
+                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
+                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)
+
+                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,
+                    )[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)

modelscope/pipelines/multi_modal/__init__.py

@@ -26,6 +26,7 @@ if TYPE_CHECKING:
     from .visual_question_answering_pipeline import VisualQuestionAnsweringPipeline
     from .video_question_answering_pipeline import VideoQuestionAnsweringPipeline
     from .videocomposer_pipeline import VideoComposerPipeline
+    from .text_to_image_freeu_pipeline import FreeUTextToImagePipeline
 else:
     _import_structure = {
         'image_captioning_pipeline': ['ImageCaptioningPipeline'],
@@ -53,7 +54,8 @@ else:
         ['SOONetVideoTemporalGroundingPipeline'],
         'text_to_video_synthesis_pipeline': ['TextToVideoSynthesisPipeline'],
         'multimodal_dialogue_pipeline': ['MultimodalDialoguePipeline'],
-        'videocomposer_pipeline': ['VideoComposerPipeline']
+        'videocomposer_pipeline': ['VideoComposerPipeline'],
+        'text_to_image_freeu_pipeline': ['FreeUTextToImagePipeline']
     }
 
     import sys

modelscope/pipelines/multi_modal/text_to_image_freeu_pipeline.py

@@ -0,0 +1,138 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os.path
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.multi_modal.freeu import (
+    register_free_crossattn_upblock2d, register_free_upblock2d)
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.base import Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+__all__ = ['FreeUTextToImagePipeline']
+
+
+@PIPELINES.register_module(
+    Tasks.text_to_image_synthesis,
+    module_name=Pipelines.freeu_stable_diffusion_text2image)
+class FreeUTextToImagePipeline(Pipeline):
+
+    def __init__(self, model=str, preprocessor=None, **kwargs):
+        """  FreeU Text to Image Pipeline.
+
+        Examples:
+
+        >>> import cv2
+        >>> from modelscope.pipelines import pipeline
+        >>> from modelscope.utils.constant import Tasks
+
+        >>> prompt = "a photo of a running corgi"  # prompt
+        >>> output_image_path = './result.png'
+        >>> inputs = {'prompt': prompt}
+        >>>
+        >>> pipe = pipeline(
+        >>>     Tasks.text_to_image_synthesis,
+        >>>     model='damo/multi-modal_freeu_stable_diffusion',
+        >>>     base_model='AI-ModelScope/stable-diffusion-v1-5',
+        >>> )
+        >>>
+        >>> output = pipe(inputs)['output_imgs']
+        >>> cv2.imwrite(output_image_path, output)
+        >>> print('pipeline: the output image path is {}'.format(output_image_path))
+        """
+        super().__init__(model=model, preprocessor=preprocessor, **kwargs)
+
+        torch_dtype = kwargs.get('torch_dtype', torch.float32)
+        self._device = getattr(
+            kwargs, 'device',
+            torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+        base_model = kwargs.get(
+            'base_model', 'AI-ModelScope/stable-diffusion-v1-5')  # default 1.5
+        self.freeu_params = kwargs.get('freeu_params', {
+            'b1': 1.5,
+            'b2': 1.6,
+            's1': 0.9,
+            's2': 0.2
+        })  # default
+
+        logger.info('load freeu stable diffusion text to image pipeline done')
+        self.pipeline = pipeline(
+            task=Tasks.text_to_image_synthesis,
+            model=base_model,
+            torch_dtype=torch_dtype,
+            device=self._device).pipeline
+
+    def preprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        return inputs
+
+    def forward(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        """
+        Inputs Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+        """
+        if not isinstance(inputs, dict):
+            raise ValueError(
+                f'Expected the input to be a dictionary, but got {type(inputs)}'
+            )
+        # -------- freeu block registration
+        register_free_upblock2d(self.pipeline, **self.freeu_params)
+        register_free_crossattn_upblock2d(self.pipeline, **self.freeu_params)
+        # -------- freeu block registration
+
+        output = self.pipeline(
+            prompt=inputs.get('prompt'),
+            height=inputs.get('height'),
+            width=inputs.get('width'),
+            num_inference_steps=inputs.get('num_inference_steps', 50),
+            guidance_scale=inputs.get('guidance_scale', 7.5),
+            negative_prompt=inputs.get('negative_prompt'),
+            num_images_per_prompt=inputs.get('num_images_per_prompt', 1),
+            eta=inputs.get('eta', 0.0),
+            generator=inputs.get('generator'),
+            latents=inputs.get('latents'),
+        ).images[0]
+
+        return {'output_tensor': output}
+
+    def postprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        output_img = np.array(inputs['output_tensor'])
+        return {OutputKeys.OUTPUT_IMGS: output_img[:, :, ::-1]}

tests/pipelines/test_text_to_image_freeu.py

@@ -0,0 +1,57 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import unittest
+
+import cv2
+
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.multi_modal import FreeUTextToImagePipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.test_utils import test_level
+
+
+class ImageEditingTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.task = Tasks.text_to_image_synthesis
+        self.model_id = 'damo/multi-modal_freeu_stable_diffusion'
+        prompt = 'a photo of a running corgi'  # prompt
+        self.inputs = {'prompt': prompt}
+        self.output_image_path = './result.png'
+        self.base_model = 'AI-ModelScope/stable-diffusion-v2-1'
+        self.freeu_params = {
+            'b1': 1.4,
+            'b2': 1.6,
+            's1': 0.9,
+            's2': 0.2
+        }  # for SD2.1
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_run_by_direct_model_download(self):
+        cache_path = snapshot_download(self.model_id)
+        pipeline = FreeUTextToImagePipeline(cache_path)
+        pipeline.group_key = self.task
+        synthesized_img = pipeline(
+            input=self.inputs)[OutputKeys.OUTPUT_IMGS]  # BGR
+        cv2.imwrite(self.output_image_path, synthesized_img)
+        print('FreeU pipeline: the synthesized image path is {}'.format(
+            self.output_image_path))
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_run_with_model_name(self):
+        pipeline_ins = pipeline(
+            task=Tasks.text_to_image_synthesis,
+            model=self.model_id,
+            base_model=self.base_model,
+            freeu_params=self.freeu_params)
+        synthesized_img = pipeline_ins(
+            self.inputs)[OutputKeys.OUTPUT_IMGS]  # BGR
+        cv2.imwrite(self.output_image_path, synthesized_img)
+        print('FreeU pipeline: the synthesized image path is {}'.format(
+            self.output_image_path))
+
+
+if __name__ == '__main__':
+    unittest.main()