diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 18d01fc0..5fc827c7 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -358,6 +358,7 @@ class Pipelines(object): ofa_text2sql = 'ofa-text2sql' video_captioning = 'video-captioning' video_question_answering = 'video-question-answering' + diffusers_stable_diffusion = 'diffusers-stable-diffusion' document_vl_embedding = 'document-vl-embedding' # science tasks diff --git a/modelscope/pipelines/multi_modal/__init__.py b/modelscope/pipelines/multi_modal/__init__.py index 07e27e14..53201a47 100644 --- a/modelscope/pipelines/multi_modal/__init__.py +++ b/modelscope/pipelines/multi_modal/__init__.py @@ -17,6 +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 else: _import_structure = { 'image_captioning_pipeline': ['ImageCaptioningPipeline'], @@ -34,7 +35,8 @@ else: 'document_vl_embedding_pipeline': ['DocumentVLEmbeddingPipeline'], 'video_captioning_pipeline': ['VideoCaptioningPipeline'], 'video_question_answering_pipeline': - ['VideoQuestionAnsweringPipeline'] + ['VideoQuestionAnsweringPipeline'], + 'diffusers_wrapped': ['StableDiffusionWrapperPipeline'] } import sys diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py new file mode 100644 index 00000000..3445a6b5 --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .stable_diffusion import StableDiffusionWrapperPipeline +else: + _import_structure = { + 'stable_diffusion': ['StableDiffusionWrapperPipeline'], + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py new file mode 100644 index 00000000..3727a460 --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/diffusers_pipeline.py @@ -0,0 +1,45 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os +from typing import Any, Dict, Generator, List, Union + +from modelscope.pipelines.base import Input, Pipeline +from modelscope.utils.constant import Hubs +from modelscope.utils.device import create_device +from modelscope.utils.hub import snapshot_download + + +class DiffusersPipeline(Pipeline): + + def __init__(self, model: str, device: str = 'gpu', **kwargs): + """ + use `model` to create a diffusers pipeline + Args: + model: model id on modelscope hub. + device: str = 'gpu' + """ + + self.device_name = device + self.cfg = None + self.preprocessor = None + self.framework = None + self.device = create_device(self.device_name) + self.hubs = kwargs.get('hubs', Hubs.modelscope) + + # make sure we download the model from modelscope hub + model_folder = model + if not os.path.isdir(model_folder): + if self.hubs != Hubs.modelscope: + raise NotImplementedError( + 'Only support model retrieval from ModelScope hub for now.' + ) + model_folder = snapshot_download(model) + + self.model = model_folder + self.models = [self.model] + self.has_multiple_models = len(self.models) > 1 + + def __call__(self, input: Union[Input, List[Input]], *args, + **kwargs) -> Union[Dict[str, Any], Generator]: + + return self.forward(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 new file mode 100644 index 00000000..55fecc8e --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .stable_diffusion_pipeline import StableDiffusionWrapperPipeline +else: + _import_structure = { + 'stable_diffusion_pipeline': ['StableDiffusionWrapperPipeline'] + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) 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 new file mode 100644 index 00000000..ef7288a9 --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py @@ -0,0 +1,46 @@ +# Copyright © Alibaba, Inc. and its affiliates. + +from typing import Any, Dict + +import torch +from diffusers import StableDiffusionPipeline + +from modelscope.metainfo import Pipelines +from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.multi_modal.diffusers_wrapped.diffusers_pipeline import \ + DiffusersPipeline +from modelscope.utils.constant import Tasks + + +# Wrap around the diffusers stable diffusion pipeline implementation +# for a unified ModelScope pipeline experience. Native stable diffusion +# pipelines will be implemented in later releases. +@PIPELINES.register_module( + Tasks.diffusers_stable_diffusion, + module_name=Pipelines.diffusers_stable_diffusion) +class StableDiffusionWrapperPipeline(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) + + # build upon the diffuser stable diffusion pipeline + self.diffusers_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 + + def forward(self, prompt, **kwargs): + return self.diffusers_pipeline(prompt, **kwargs) + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 108dcedc..f73ec428 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -185,6 +185,7 @@ 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 822ac43f..451db10f 100644 --- a/requirements/multi-modal.txt +++ b/requirements/multi-modal.txt @@ -1,3 +1,4 @@ +diffusers>=0.11.1 ftfy>=6.0.3 librosa opencv-python diff --git a/tests/pipelines/test_diffusers_stable_diffusion.py b/tests/pipelines/test_diffusers_stable_diffusion.py new file mode 100644 index 00000000..7784bff6 --- /dev/null +++ b/tests/pipelines/test_diffusers_stable_diffusion.py @@ -0,0 +1,28 @@ +# 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 DiffusersStableDiffusionTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.diffusers_stable_diffusion + self.model_id = 'shadescript/stable-diffusion-2-1-dev' + + test_input = 'a photo of an astronaut riding a horse on mars' + + @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.images[0].save('/tmp/output.png') + print('Image saved to /tmp/output.png') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/pipelines/test_video_depth_estimation.py b/tests/pipelines/test_video_depth_estimation.py index 2fd0b872..77cb4b9b 100644 --- a/tests/pipelines/test_video_depth_estimation.py +++ b/tests/pipelines/test_video_depth_estimation.py @@ -15,7 +15,7 @@ class VideoDepthEstimationTest(unittest.TestCase, DemoCompatibilityCheck): self.task = 'video-depth-estimation' self.model_id = 'damo/cv_dro-resnet18_video-depth-estimation_indoor' - @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 3, 'skip test in current test level') def test_image_depth_estimation(self): input_location = 'data/test/videos/video_depth_estimation.mp4' estimator = pipeline(Tasks.video_depth_estimation, model=self.model_id)