Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
shivamshrirao
GitHub Repository: shivamshrirao/diffusers
 Path: blob/main/tests/pipelines/stable_diffusion/test_stable_diffusion_k_diffusion.py
1451 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 gc

17
import unittest

18


19
import numpy as np

20
import torch

21


22
from diffusers import StableDiffusionKDiffusionPipeline

23
from diffusers.utils import slow, torch_device

24
from diffusers.utils.testing_utils import require_torch_gpu

25


26


27
torch.backends.cuda.matmul.allow_tf32 = False

28


29


30
@slow

31
@require_torch_gpu

32
class StableDiffusionPipelineIntegrationTests(unittest.TestCase):

33
    def tearDown(self):

34
        # clean up the VRAM after each test

35
        super().tearDown()

36
        gc.collect()

37
        torch.cuda.empty_cache()

38


39
    def test_stable_diffusion_1(self):

40
        sd_pipe = StableDiffusionKDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

41
        sd_pipe = sd_pipe.to(torch_device)

42
        sd_pipe.set_progress_bar_config(disable=None)

43


44
        sd_pipe.set_scheduler("sample_euler")

45


46
        prompt = "A painting of a squirrel eating a burger"

47
        generator = torch.manual_seed(0)

48
        output = sd_pipe([prompt], generator=generator, guidance_scale=9.0, num_inference_steps=20, output_type="np")

49


50
        image = output.images

51


52
        image_slice = image[0, -3:, -3:, -1]

53


54
        assert image.shape == (1, 512, 512, 3)

55
        expected_slice = np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339])

56


57
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2

58


59
    def test_stable_diffusion_2(self):

60
        sd_pipe = StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")

61
        sd_pipe = sd_pipe.to(torch_device)

62
        sd_pipe.set_progress_bar_config(disable=None)

63


64
        sd_pipe.set_scheduler("sample_euler")

65


66
        prompt = "A painting of a squirrel eating a burger"

67
        generator = torch.manual_seed(0)

68
        output = sd_pipe([prompt], generator=generator, guidance_scale=9.0, num_inference_steps=20, output_type="np")

69


70
        image = output.images

71


72
        image_slice = image[0, -3:, -3:, -1]

73


74
        assert image.shape == (1, 512, 512, 3)

75
        expected_slice = np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112])

76


77
        assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-1

78


79