diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index ea56efb5..377ade9b 100644 --- a/modelscope/metainfo.py +++ b/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' diff --git a/modelscope/models/multi_modal/freeu/__init__.py b/modelscope/models/multi_modal/freeu/__init__.py new file mode 100644 index 00000000..3cd55cf3 --- /dev/null +++ b/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={}, + ) diff --git a/modelscope/models/multi_modal/freeu/free_lunch_utils.py b/modelscope/models/multi_modal/freeu/free_lunch_utils.py new file mode 100644 index 00000000..eb5d191f --- /dev/null +++ b/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) diff --git a/modelscope/pipelines/multi_modal/__init__.py b/modelscope/pipelines/multi_modal/__init__.py index b5316684..1faa261e 100644 --- a/modelscope/pipelines/multi_modal/__init__.py +++ b/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 diff --git a/modelscope/pipelines/multi_modal/text_to_image_freeu_pipeline.py b/modelscope/pipelines/multi_modal/text_to_image_freeu_pipeline.py new file mode 100644 index 00000000..9300554c --- /dev/null +++ b/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]} diff --git a/tests/pipelines/test_text_to_image_freeu.py b/tests/pipelines/test_text_to_image_freeu.py new file mode 100644 index 00000000..7aebe318 --- /dev/null +++ b/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()