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
This commit is contained in:
ly103369
2023-01-11 21:49:58 +08:00
committed by wenmeng.zwm
parent b1b65c2597
commit 0e54e80b23
11 changed files with 251 additions and 17 deletions

View File

@@ -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'

View File

@@ -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

View File

@@ -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

View File

@@ -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))

View File

@@ -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

View File

@@ -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

View File

@@ -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}

View File

@@ -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):

View File

@@ -1,3 +1,4 @@
accelerate
diffusers>=0.11.1
ftfy>=6.0.3
librosa

View File

@@ -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()

View File

@@ -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__':