diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index d8b95ab4..dec5084f 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -26,6 +26,7 @@ class OutputKeys(object): POLYGONS = 'polygons' OUTPUT = 'output' OUTPUT_IMG = 'output_img' + OUTPUT_IMGS = 'output_imgs' OUTPUT_VIDEO = 'output_video' OUTPUT_PCM = 'output_pcm' OUTPUT_PCM_LIST = 'output_pcm_list' @@ -797,11 +798,11 @@ TASK_OUTPUTS = { # } Tasks.visual_grounding: [OutputKeys.BOXES, OutputKeys.SCORES], - # text_to_image result for a single sample + # text_to_image result for samples # { - # "output_img": np.ndarray with shape [height, width, 3] + # "output_imgs": np.ndarray list with shape [[height, width, 3], ...] # } - Tasks.text_to_image_synthesis: [OutputKeys.OUTPUT_IMG], + Tasks.text_to_image_synthesis: [OutputKeys.OUTPUT_IMGS], # text_to_speech result for a single sample # { diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py index efec87d1..d1e3a2ae 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py @@ -47,6 +47,11 @@ class DiffusersPipeline(Pipeline): def __call__(self, input: Union[Input, List[Input]], *args, **kwargs) -> Union[Dict[str, Any], Generator]: - - return self.postprocess( - self.forward(self.preprocess(input), *args, **kwargs)) + preprocess_params, forward_params, postprocess_params = self._sanitize_parameters( + **kwargs) + self._check_input(input) + out = self.preprocess(input, **preprocess_params) + out = self.forward(out, **forward_params) + out = self.postprocess(out, **postprocess_params) + self._check_output(out) + return out 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 index 8bf25ba3..fa8e1f50 100644 --- 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 @@ -8,6 +8,8 @@ from typing import Any, Dict, List, Union +import cv2 +import numpy as np import torch from diffusers import StableDiffusionPipeline from diffusers.models import AutoencoderKL, UNet2DConditionModel @@ -16,6 +18,7 @@ from diffusers.schedulers import (DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler) +from PIL import Image from transformers import (ChineseCLIPProcessor, ChineseCLIPTextModel, CLIPFeatureExtractor) @@ -41,15 +44,42 @@ class ChineseStableDiffusionPipeline(DiffusersPipeline): """ super().__init__(model, device, **kwargs) - torch_dtype = kwargs.get('torch_dtype', torch.float16) + torch_dtype = kwargs.get('torch_dtype', torch.float32) self.pipeline = _DiffuersChineseStableDiffusionPipeline.from_pretrained( model, torch_dtype=torch_dtype).to(self.device) - def forward(self, prompt, **kwargs): - return self.pipeline(prompt, **kwargs) + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + if not isinstance(inputs, dict): + raise ValueError( + f'Expected the input to be a dictionary, but got {type(input)}' + ) + if 'text' not in inputs: + raise ValueError('input should contain "text", but not found') + + return self.pipeline( + prompt=inputs.get('text'), + 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'), + output_type=inputs.get('output_type', 'pil'), + return_dict=inputs.get('return_dict', True), + callback=inputs.get('callback'), + callback_steps=inputs.get('callback_steps', 1)) def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return {OutputKeys.OUTPUT_IMG: inputs.images} + images = [] + for img in inputs.images: + if isinstance(img, Image.Image): + img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) + images.append(img) + return {OutputKeys.OUTPUT_IMGS: images} class _DiffuersChineseStableDiffusionPipeline(StableDiffusionPipeline): 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 0e89200f..b6d9d3bd 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 @@ -2,8 +2,11 @@ from typing import Any, Dict +import cv2 +import numpy as np import torch from diffusers import StableDiffusionPipeline +from PIL import Image from modelscope.metainfo import Pipelines from modelscope.outputs import OutputKeys @@ -30,15 +33,42 @@ class StableDiffusionWrapperPipeline(DiffusersPipeline): """ super().__init__(model, device, **kwargs) - torch_dtype = kwargs.get('torch_dtype', torch.float16) + torch_dtype = kwargs.get('torch_dtype', torch.float32) # build upon the diffuser stable diffusion pipeline self.pipeline = StableDiffusionPipeline.from_pretrained( model, torch_dtype=torch_dtype) self.pipeline.to(self.device) - def forward(self, prompt, **kwargs): - return self.pipeline(prompt, **kwargs) + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + if not isinstance(inputs, dict): + raise ValueError( + f'Expected the input to be a dictionary, but got {type(input)}' + ) + if 'text' not in inputs: + raise ValueError('input should contain "text", but not found') + + return self.pipeline( + prompt=inputs.get('text'), + 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'), + output_type=inputs.get('output_type', 'pil'), + return_dict=inputs.get('return_dict', True), + callback=inputs.get('callback'), + callback_steps=inputs.get('callback_steps', 1)) def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return {OutputKeys.OUTPUT_IMG: inputs.images} + images = [] + for img in inputs.images: + if isinstance(img, Image.Image): + img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) + images.append(img) + return {OutputKeys.OUTPUT_IMGS: images} diff --git a/modelscope/pipelines/multi_modal/text_to_image_synthesis_pipeline.py b/modelscope/pipelines/multi_modal/text_to_image_synthesis_pipeline.py index 36e761aa..42e9c819 100644 --- a/modelscope/pipelines/multi_modal/text_to_image_synthesis_pipeline.py +++ b/modelscope/pipelines/multi_modal/text_to_image_synthesis_pipeline.py @@ -50,4 +50,6 @@ class TextToImageSynthesisPipeline(Pipeline): return self.model.generate(input) def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return {OutputKeys.OUTPUT_IMG: inputs} + if not isinstance(inputs, list): + inputs = [inputs] + return {OutputKeys.OUTPUT_IMGS: inputs} diff --git a/tests/pipelines/test_chinese_stable_diffusion.py b/tests/pipelines/test_chinese_stable_diffusion.py index acbdb074..bd6d74aa 100644 --- a/tests/pipelines/test_chinese_stable_diffusion.py +++ b/tests/pipelines/test_chinese_stable_diffusion.py @@ -2,6 +2,8 @@ import unittest +import cv2 + from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from modelscope.utils.demo_utils import DemoCompatibilityCheck @@ -17,8 +19,8 @@ class ChineseStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): @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') + output = pipe({'text': '中国山水画'}) + cv2.imwrite('result.png', output['output_imgs'][0]) print('Image saved to result.png') @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') @@ -27,8 +29,8 @@ class ChineseStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): 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') + output = pipe({'text': '中国山水画', 'num_inference_steps': 25}) + cv2.imwrite('result2.png', output['output_imgs'][0]) print('Image saved to result2.png') diff --git a/tests/pipelines/test_diffusers_stable_diffusion.py b/tests/pipelines/test_diffusers_stable_diffusion.py index 4ffc4d26..98c4862a 100644 --- a/tests/pipelines/test_diffusers_stable_diffusion.py +++ b/tests/pipelines/test_diffusers_stable_diffusion.py @@ -2,6 +2,8 @@ import unittest +import cv2 + from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from modelscope.utils.demo_utils import DemoCompatibilityCheck @@ -19,8 +21,12 @@ class DiffusersStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') 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['output_img'][0].save('output.png') + output = diffusers_pipeline({ + 'text': self.test_input, + 'height': 512, + 'width': 512 + }) + cv2.imwrite('output.png', output['output_imgs'][0]) print('Image saved to output.png') diff --git a/tests/pipelines/test_multi_stage_diffusion.py b/tests/pipelines/test_multi_stage_diffusion.py index f4e63ce0..802c8ed7 100644 --- a/tests/pipelines/test_multi_stage_diffusion.py +++ b/tests/pipelines/test_multi_stage_diffusion.py @@ -23,7 +23,7 @@ class MultiStageDiffusionTest(unittest.TestCase): pipe_line_text_to_image_synthesis = pipeline( task=Tasks.text_to_image_synthesis, model=model) img = pipe_line_text_to_image_synthesis( - self.test_text)[OutputKeys.OUTPUT_IMG] + self.test_text)[OutputKeys.OUTPUT_IMGS][0] print(np.sum(np.abs(img))) @unittest.skip( @@ -32,7 +32,7 @@ class MultiStageDiffusionTest(unittest.TestCase): pipe_line_text_to_image_synthesis = pipeline( task=Tasks.text_to_image_synthesis, model=self.model_id) img = pipe_line_text_to_image_synthesis( - self.test_text)[OutputKeys.OUTPUT_IMG] + self.test_text)[OutputKeys.OUTPUT_IMGS][0] print(np.sum(np.abs(img))) diff --git a/tests/pipelines/test_ofa_tasks.py b/tests/pipelines/test_ofa_tasks.py index 8dc7197d..df1b5647 100644 --- a/tests/pipelines/test_ofa_tasks.py +++ b/tests/pipelines/test_ofa_tasks.py @@ -303,7 +303,7 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): ofa_pipe.model.generator.beam_size = 2 example = {'text': 'a bear in the water.'} result = ofa_pipe(example) - result[OutputKeys.OUTPUT_IMG].save('result.png') + result[OutputKeys.OUTPUT_IMGS][0].save('result.png') print(f'Output written to {osp.abspath("result.png")}') @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') diff --git a/tests/pipelines/test_text_to_image_synthesis.py b/tests/pipelines/test_text_to_image_synthesis.py index 82f4e657..5e28282b 100644 --- a/tests/pipelines/test_text_to_image_synthesis.py +++ b/tests/pipelines/test_text_to_image_synthesis.py @@ -32,7 +32,7 @@ class TextToImageSynthesisTest(unittest.TestCase, DemoCompatibilityCheck): pipe_line_text_to_image_synthesis = pipeline( task=Tasks.text_to_image_synthesis, model=model) img = pipe_line_text_to_image_synthesis( - self.test_text)[OutputKeys.OUTPUT_IMG] + self.test_text)[OutputKeys.OUTPUT_IMGS][0] print(np.sum(np.abs(img))) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') @@ -40,7 +40,7 @@ class TextToImageSynthesisTest(unittest.TestCase, DemoCompatibilityCheck): pipe_line_text_to_image_synthesis = pipeline( task=Tasks.text_to_image_synthesis, model=self.model_id) img = pipe_line_text_to_image_synthesis( - self.test_text)[OutputKeys.OUTPUT_IMG] + self.test_text)[OutputKeys.OUTPUT_IMGS][0] print(np.sum(np.abs(img))) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') @@ -48,7 +48,7 @@ class TextToImageSynthesisTest(unittest.TestCase, DemoCompatibilityCheck): pipe_line_text_to_image_synthesis = pipeline( task=Tasks.text_to_image_synthesis) img = pipe_line_text_to_image_synthesis( - self.test_text)[OutputKeys.OUTPUT_IMG] + self.test_text)[OutputKeys.OUTPUT_IMGS][0] print(np.sum(np.abs(img))) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') @@ -58,7 +58,7 @@ class TextToImageSynthesisTest(unittest.TestCase, DemoCompatibilityCheck): pipe_line_text_to_image_synthesis = pipeline( task=Tasks.text_to_image_synthesis, model=model) img = pipe_line_text_to_image_synthesis( - self.test_text)[OutputKeys.OUTPUT_IMG] + self.test_text)[OutputKeys.OUTPUT_IMGS][0] print(np.sum(np.abs(img))) @unittest.skip('demo compatibility test is only enabled on a needed-basis')