CoCalc -- test_onnx_stable_diffusion

GitHub Repository: shivamshrirao/diffusers
Path: blob/main/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py
¹⁴⁴⁸ views
1
# coding=utf-8
2
# Copyright 2023 HuggingFace Inc.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
import random
17
import unittest
18

19
import numpy as np
20

21
from diffusers import (
22
    DPMSolverMultistepScheduler,
23
    EulerAncestralDiscreteScheduler,
24
    EulerDiscreteScheduler,
25
    LMSDiscreteScheduler,
26
    OnnxStableDiffusionImg2ImgPipeline,
27
    PNDMScheduler,
28
)
29
from diffusers.utils import floats_tensor
30
from diffusers.utils.testing_utils import (
31
    is_onnx_available,
32
    load_image,
33
    nightly,
34
    require_onnxruntime,
35
    require_torch_gpu,
36
)
37

38
from ...test_pipelines_onnx_common import OnnxPipelineTesterMixin
39

40

41
if is_onnx_available():
42
    import onnxruntime as ort
43

44

45
class OnnxStableDiffusionImg2ImgPipelineFastTests(OnnxPipelineTesterMixin, unittest.TestCase):
46
    hub_checkpoint = "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"
47

48
    def get_dummy_inputs(self, seed=0):
49
        image = floats_tensor((1, 3, 128, 128), rng=random.Random(seed))
50
        generator = np.random.RandomState(seed)
51
        inputs = {
52
            "prompt": "A painting of a squirrel eating a burger",
53
            "image": image,
54
            "generator": generator,
55
            "num_inference_steps": 3,
56
            "strength": 0.75,
57
            "guidance_scale": 7.5,
58
            "output_type": "numpy",
59
        }
60
        return inputs
61

62
    def test_pipeline_default_ddim(self):
63
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
64
        pipe.set_progress_bar_config(disable=None)
65

66
        inputs = self.get_dummy_inputs()
67
        image = pipe(**inputs).images
68
        image_slice = image[0, -3:, -3:, -1].flatten()
69

70
        assert image.shape == (1, 128, 128, 3)
71
        expected_slice = np.array([0.69643, 0.58484, 0.50314, 0.58760, 0.55368, 0.59643, 0.51529, 0.41217, 0.49087])
72
        assert np.abs(image_slice - expected_slice).max() < 1e-1
73

74
    def test_pipeline_pndm(self):
75
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
76
        pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config, skip_prk_steps=True)
77
        pipe.set_progress_bar_config(disable=None)
78

79
        inputs = self.get_dummy_inputs()
80
        image = pipe(**inputs).images
81
        image_slice = image[0, -3:, -3:, -1]
82

83
        assert image.shape == (1, 128, 128, 3)
84
        expected_slice = np.array([0.61710, 0.53390, 0.49310, 0.55622, 0.50982, 0.58240, 0.50716, 0.38629, 0.46856])
85

86
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1
87

88
    def test_pipeline_lms(self):
89
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
90
        pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
91
        pipe.set_progress_bar_config(disable=None)
92

93
        # warmup pass to apply optimizations
94
        _ = pipe(**self.get_dummy_inputs())
95

96
        inputs = self.get_dummy_inputs()
97
        image = pipe(**inputs).images
98
        image_slice = image[0, -3:, -3:, -1]
99

100
        assert image.shape == (1, 128, 128, 3)
101
        expected_slice = np.array([0.52761, 0.59977, 0.49033, 0.49619, 0.54282, 0.50311, 0.47600, 0.40918, 0.45203])
102

103
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1
104

105
    def test_pipeline_euler(self):
106
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
107
        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
108
        pipe.set_progress_bar_config(disable=None)
109

110
        inputs = self.get_dummy_inputs()
111
        image = pipe(**inputs).images
112
        image_slice = image[0, -3:, -3:, -1]
113

114
        assert image.shape == (1, 128, 128, 3)
115
        expected_slice = np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304])
116

117
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1
118

119
    def test_pipeline_euler_ancestral(self):
120
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
121
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
122
        pipe.set_progress_bar_config(disable=None)
123

124
        inputs = self.get_dummy_inputs()
125
        image = pipe(**inputs).images
126
        image_slice = image[0, -3:, -3:, -1]
127

128
        assert image.shape == (1, 128, 128, 3)
129
        expected_slice = np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304])
130

131
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1
132

133
    def test_pipeline_dpm_multistep(self):
134
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
135
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
136
        pipe.set_progress_bar_config(disable=None)
137

138
        inputs = self.get_dummy_inputs()
139
        image = pipe(**inputs).images
140
        image_slice = image[0, -3:, -3:, -1]
141

142
        assert image.shape == (1, 128, 128, 3)
143
        expected_slice = np.array([0.65331, 0.58277, 0.48204, 0.56059, 0.53665, 0.56235, 0.50969, 0.40009, 0.46552])
144

145
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1
146

147

148
@nightly
149
@require_onnxruntime
150
@require_torch_gpu
151
class OnnxStableDiffusionImg2ImgPipelineIntegrationTests(unittest.TestCase):
152
    @property
153
    def gpu_provider(self):
154
        return (
155
            "CUDAExecutionProvider",
156
            {
157
                "gpu_mem_limit": "15000000000",  # 15GB
158
                "arena_extend_strategy": "kSameAsRequested",
159
            },
160
        )
161

162
    @property
163
    def gpu_options(self):
164
        options = ort.SessionOptions()
165
        options.enable_mem_pattern = False
166
        return options
167

168
    def test_inference_default_pndm(self):
169
        init_image = load_image(
170
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
171
            "/img2img/sketch-mountains-input.jpg"
172
        )
173
        init_image = init_image.resize((768, 512))
174
        # using the PNDM scheduler by default
175
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(
176
            "CompVis/stable-diffusion-v1-4",
177
            revision="onnx",
178
            safety_checker=None,
179
            feature_extractor=None,
180
            provider=self.gpu_provider,
181
            sess_options=self.gpu_options,
182
        )
183
        pipe.set_progress_bar_config(disable=None)
184

185
        prompt = "A fantasy landscape, trending on artstation"
186

187
        generator = np.random.RandomState(0)
188
        output = pipe(
189
            prompt=prompt,
190
            image=init_image,
191
            strength=0.75,
192
            guidance_scale=7.5,
193
            num_inference_steps=10,
194
            generator=generator,
195
            output_type="np",
196
        )
197
        images = output.images
198
        image_slice = images[0, 255:258, 383:386, -1]
199

200
        assert images.shape == (1, 512, 768, 3)
201
        expected_slice = np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019])
202
        # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues
203

204
        assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2
205

206
    def test_inference_k_lms(self):
207
        init_image = load_image(
208
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
209
            "/img2img/sketch-mountains-input.jpg"
210
        )
211
        init_image = init_image.resize((768, 512))
212
        lms_scheduler = LMSDiscreteScheduler.from_pretrained(
213
            "runwayml/stable-diffusion-v1-5", subfolder="scheduler", revision="onnx"
214
        )
215
        pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(
216
            "runwayml/stable-diffusion-v1-5",
217
            revision="onnx",
218
            scheduler=lms_scheduler,
219
            safety_checker=None,
220
            feature_extractor=None,
221
            provider=self.gpu_provider,
222
            sess_options=self.gpu_options,
223
        )
224
        pipe.set_progress_bar_config(disable=None)
225

226
        prompt = "A fantasy landscape, trending on artstation"
227

228
        generator = np.random.RandomState(0)
229
        output = pipe(
230
            prompt=prompt,
231
            image=init_image,
232
            strength=0.75,
233
            guidance_scale=7.5,
234
            num_inference_steps=20,
235
            generator=generator,
236
            output_type="np",
237
        )
238
        images = output.images
239
        image_slice = images[0, 255:258, 383:386, -1]
240

241
        assert images.shape == (1, 512, 768, 3)
242
        expected_slice = np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431])
243
        # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues
244

245
        assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2
246

247
Product

Resources

Company
