From 0e54e80b23850badcc1b5542c6eebf3925935b04 Mon Sep 17 00:00:00 2001 From: ly103369 Date: Wed, 11 Jan 2023 21:49:58 +0800 Subject: [PATCH] add Damo Chinese Stable Diffusion and fix bugs in DiffusersPipeline Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11386214 * chinese sd & fix sdwrapper * impl preprocess and postprocess in DiffusersPipeline * add accelerate requirements --- modelscope/metainfo.py | 1 + modelscope/pipelines/multi_modal/__init__.py | 5 +- .../multi_modal/diffusers_wrapped/__init__.py | 4 +- .../diffusers_wrapped/diffusers_pipeline.py | 9 +- .../stable_diffusion/__init__.py | 5 +- .../chinese_stable_diffusion_pipeline.py | 186 ++++++++++++++++++ .../stable_diffusion_pipeline.py | 14 +- modelscope/utils/constant.py | 1 - requirements/multi-modal.txt | 1 + .../test_chinese_stable_diffusion.py | 36 ++++ .../test_diffusers_stable_diffusion.py | 6 +- 11 files changed, 251 insertions(+), 17 deletions(-) create mode 100644 modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/chinese_stable_diffusion_pipeline.py create mode 100644 tests/pipelines/test_chinese_stable_diffusion.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 782692b7..8bb95cbb 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -367,6 +367,7 @@ class Pipelines(object): video_question_answering = 'video-question-answering' diffusers_stable_diffusion = 'diffusers-stable-diffusion' document_vl_embedding = 'document-vl-embedding' + chinese_stable_diffusion = 'chinese-stable-diffusion' # science tasks protein_structure = 'unifold-protein-structure' diff --git a/modelscope/pipelines/multi_modal/__init__.py b/modelscope/pipelines/multi_modal/__init__.py index 53201a47..8d914aee 100644 --- a/modelscope/pipelines/multi_modal/__init__.py +++ b/modelscope/pipelines/multi_modal/__init__.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: from .document_vl_embedding_pipeline import DocumentVLEmbeddingPipeline from .video_captioning_pipeline import VideoCaptioningPipeline from .video_question_answering_pipeline import VideoQuestionAnsweringPipeline - from .diffusers_wrapped import StableDiffusionWrapperPipeline + from .diffusers_wrapped import StableDiffusionWrapperPipeline, ChineseStableDiffusionPipeline else: _import_structure = { 'image_captioning_pipeline': ['ImageCaptioningPipeline'], @@ -36,7 +36,8 @@ else: 'video_captioning_pipeline': ['VideoCaptioningPipeline'], 'video_question_answering_pipeline': ['VideoQuestionAnsweringPipeline'], - 'diffusers_wrapped': ['StableDiffusionWrapperPipeline'] + 'diffusers_wrapped': + ['StableDiffusionWrapperPipeline', 'ChineseStableDiffusionPipeline'] } import sys diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py index 3445a6b5..0c9fc5e8 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py @@ -5,9 +5,11 @@ from modelscope.utils.import_utils import LazyImportModule if TYPE_CHECKING: from .stable_diffusion import StableDiffusionWrapperPipeline + from .stable_diffusion import ChineseStableDiffusionPipeline else: _import_structure = { - 'stable_diffusion': ['StableDiffusionWrapperPipeline'], + 'stable_diffusion': + ['StableDiffusionWrapperPipeline', 'ChineseStableDiffusionPipeline'] } import sys diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py index 3727a460..efec87d1 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py @@ -39,7 +39,14 @@ class DiffusersPipeline(Pipeline): self.models = [self.model] self.has_multiple_models = len(self.models) > 1 + def preprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs + def __call__(self, input: Union[Input, List[Input]], *args, **kwargs) -> Union[Dict[str, Any], Generator]: - return self.forward(input, *args, **kwargs) + return self.postprocess( + self.forward(self.preprocess(input), *args, **kwargs)) diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py index 55fecc8e..6892877a 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py @@ -5,9 +5,12 @@ from modelscope.utils.import_utils import LazyImportModule if TYPE_CHECKING: from .stable_diffusion_pipeline import StableDiffusionWrapperPipeline + from .chinese_stable_diffusion_pipeline import ChineseStableDiffusionPipeline else: _import_structure = { - 'stable_diffusion_pipeline': ['StableDiffusionWrapperPipeline'] + 'stable_diffusion_pipeline': ['StableDiffusionWrapperPipeline'], + 'chinese_stable_diffusion_pipeline': + ['ChineseStableDiffusionPipeline'] } import sys diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/chinese_stable_diffusion_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/chinese_stable_diffusion_pipeline.py new file mode 100644 index 00000000..8bf25ba3 --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/chinese_stable_diffusion_pipeline.py @@ -0,0 +1,186 @@ +# Copyright 2022 The HuggingFace Team. +# Copyright 2022 The Alibaba Fundamental Vision Team Authors. All rights reserved. + +# The implementation here is modified based on diffusers, +# originally Apache License, Copyright 2022 The HuggingFace Team, +# and publicly available at +# https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py + +from typing import Any, Dict, List, Union + +import torch +from diffusers import StableDiffusionPipeline +from diffusers.models import AutoencoderKL, UNet2DConditionModel +from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker +from diffusers.schedulers import (DDIMScheduler, DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, LMSDiscreteScheduler, + PNDMScheduler) +from transformers import (ChineseCLIPProcessor, ChineseCLIPTextModel, + CLIPFeatureExtractor) + +from modelscope.metainfo import Pipelines +from modelscope.outputs import OutputKeys +from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.multi_modal.diffusers_wrapped.diffusers_pipeline import \ + DiffusersPipeline +from modelscope.utils.constant import Tasks + + +@PIPELINES.register_module( + Tasks.text_to_image_synthesis, + module_name=Pipelines.chinese_stable_diffusion) +class ChineseStableDiffusionPipeline(DiffusersPipeline): + + def __init__(self, model: str, device: str = 'gpu', **kwargs): + """ + use `model` to create a stable diffusion pipeline + Args: + model: model id on modelscope hub. + device: str = 'gpu' + """ + super().__init__(model, device, **kwargs) + + torch_dtype = kwargs.get('torch_dtype', torch.float16) + self.pipeline = _DiffuersChineseStableDiffusionPipeline.from_pretrained( + model, torch_dtype=torch_dtype).to(self.device) + + def forward(self, prompt, **kwargs): + return self.pipeline(prompt, **kwargs) + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return {OutputKeys.OUTPUT_IMG: inputs.images} + + +class _DiffuersChineseStableDiffusionPipeline(StableDiffusionPipeline): + r""" + Pipeline for text-to-image generation using Chinese Stable Diffusion. + + This model inherits from [`StableDiffusionPipeline`]. Check the superclass documentation for the generic methods the + library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`ChineseCLIPTextModel`]): + Frozen text-encoder. Chinese Stable Diffusion uses the text portion of [ChineseCLIP] + (https://huggingface.co/docs/transformers/main/en/model_doc/chinese_clip#transformers.ChineseCLIPTextModel), + specifically the [chinese-clip-vit-huge-patch14] + (https://huggingface.co/OFA-Sys/chinese-clip-vit-huge-patch14) variant. + tokenizer (`ChineseCLIPProcessor`): + Tokenizer of class + [ChineseCLIPProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/chinese_clip#transformers.ChineseCLIPProcessor). + unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + safety_checker ([`StableDiffusionSafetyChecker`]): + Classification module that estimates whether generated images could be considered offensive or harmful. + Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + feature_extractor ([`CLIPFeatureExtractor`]): + Model that extracts features from generated images to be used as inputs for the `safety_checker`. + """ + _optional_components = ['safety_checker', 'feature_extractor'] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: ChineseCLIPTextModel, + tokenizer: ChineseCLIPProcessor, + unet: UNet2DConditionModel, + scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, ], + safety_checker: StableDiffusionSafetyChecker, + feature_extractor: CLIPFeatureExtractor, + requires_safety_checker: bool = True, + ): + super().__init__( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + safety_checker=safety_checker, + feature_extractor=feature_extractor, + requires_safety_checker=requires_safety_checker) + + def _encode_prompt(self, prompt, device, num_images_per_prompt, + do_classifier_free_guidance, negative_prompt): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list(int)`): + prompt to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`): + The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored + if `guidance_scale` is less than `1`). + """ + batch_size = len(prompt) if isinstance(prompt, list) else 1 + + text_inputs = self.tokenizer( + text=prompt, + padding='max_length', + truncation=True, + max_length=52, + return_tensors='pt') + text_inputs = {k: v.to(device) for k, v in text_inputs.items()} + text_embeddings = self.text_encoder(**text_inputs) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1) + text_embeddings = text_embeddings.view( + bs_embed * num_images_per_prompt, seq_len, -1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [''] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f'`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=' + f' {type(prompt)}.') + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f'`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:' + f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches' + ' the batch size of `prompt`.') + else: + uncond_tokens = negative_prompt + + uncond_input = self.tokenizer( + text=uncond_tokens, + padding='max_length', + truncation=True, + max_length=52, + return_tensors='pt') + uncond_input = {k: v.to(device) for k, v in uncond_input.items()} + uncond_embeddings = self.text_encoder(**uncond_input) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + uncond_embeddings = uncond_embeddings.repeat( + 1, num_images_per_prompt, 1) + uncond_embeddings = uncond_embeddings.view( + batch_size * num_images_per_prompt, seq_len, -1) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py index ef7288a9..0e89200f 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py @@ -6,6 +6,7 @@ import torch from diffusers import StableDiffusionPipeline from modelscope.metainfo import Pipelines +from modelscope.outputs import OutputKeys from modelscope.pipelines.builder import PIPELINES from modelscope.pipelines.multi_modal.diffusers_wrapped.diffusers_pipeline import \ DiffusersPipeline @@ -16,7 +17,7 @@ from modelscope.utils.constant import Tasks # for a unified ModelScope pipeline experience. Native stable diffusion # pipelines will be implemented in later releases. @PIPELINES.register_module( - Tasks.diffusers_stable_diffusion, + Tasks.text_to_image_synthesis, module_name=Pipelines.diffusers_stable_diffusion) class StableDiffusionWrapperPipeline(DiffusersPipeline): @@ -32,15 +33,12 @@ class StableDiffusionWrapperPipeline(DiffusersPipeline): torch_dtype = kwargs.get('torch_dtype', torch.float16) # build upon the diffuser stable diffusion pipeline - self.diffusers_pipeline = StableDiffusionPipeline.from_pretrained( + self.pipeline = StableDiffusionPipeline.from_pretrained( model, torch_dtype=torch_dtype) - self.diffusers_pipeline.to(self.device) - - def preprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return inputs + self.pipeline.to(self.device) def forward(self, prompt, **kwargs): - return self.diffusers_pipeline(prompt, **kwargs) + return self.pipeline(prompt, **kwargs) def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return inputs + return {OutputKeys.OUTPUT_IMG: inputs.images} diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 2b313680..26b68587 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -187,7 +187,6 @@ class MultiModalTasks(object): document_vl_embedding = 'document-vl-embedding' video_captioning = 'video-captioning' video_question_answering = 'video-question-answering' - diffusers_stable_diffusion = 'diffusers-stable-diffusion' class ScienceTasks(object): diff --git a/requirements/multi-modal.txt b/requirements/multi-modal.txt index 451db10f..8a86be8e 100644 --- a/requirements/multi-modal.txt +++ b/requirements/multi-modal.txt @@ -1,3 +1,4 @@ +accelerate diffusers>=0.11.1 ftfy>=6.0.3 librosa diff --git a/tests/pipelines/test_chinese_stable_diffusion.py b/tests/pipelines/test_chinese_stable_diffusion.py new file mode 100644 index 00000000..acbdb074 --- /dev/null +++ b/tests/pipelines/test_chinese_stable_diffusion.py @@ -0,0 +1,36 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import unittest + +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class ChineseStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.text_to_image_synthesis + self.model_id = 'damo/multi-modal_chinese_stable_diffusion_v1.0' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_default(self): + pipe = pipeline(task=self.task, model=self.model_id) + output = pipe('中国山水画') + output['output_img'][0].save('result.png') + print('Image saved to result.png') + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_dpmsolver(self): + from diffusers.schedulers import DPMSolverMultistepScheduler + pipe = pipeline(task=self.task, model=self.model_id) + pipe.pipeline.scheduler = DPMSolverMultistepScheduler.from_config( + pipe.pipeline.scheduler.config) + output = pipe('中国山水画') + output['output_img'][0].save('result2.png') + print('Image saved to result2.png') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/pipelines/test_diffusers_stable_diffusion.py b/tests/pipelines/test_diffusers_stable_diffusion.py index 7784bff6..4ffc4d26 100644 --- a/tests/pipelines/test_diffusers_stable_diffusion.py +++ b/tests/pipelines/test_diffusers_stable_diffusion.py @@ -11,7 +11,7 @@ from modelscope.utils.test_utils import test_level class DiffusersStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): def setUp(self) -> None: - self.task = Tasks.diffusers_stable_diffusion + self.task = Tasks.text_to_image_synthesis self.model_id = 'shadescript/stable-diffusion-2-1-dev' test_input = 'a photo of an astronaut riding a horse on mars' @@ -20,8 +20,8 @@ class DiffusersStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): def test_run(self): diffusers_pipeline = pipeline(task=self.task, model=self.model_id) output = diffusers_pipeline(self.test_input, height=512, width=512) - output.images[0].save('/tmp/output.png') - print('Image saved to /tmp/output.png') + output['output_img'][0].save('output.png') + print('Image saved to output.png') if __name__ == '__main__':