




import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


snake_case_ : Optional[Any]    		=       False







class __snake_case			(			unittest.TestCase		):
   pass







@nightly
@require_torch_gpu
class __snake_case			(			unittest.TestCase		):







   def lowerCamelCase						(      self		:						int):

        """simple docstring"""


        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


        UpperCAmelCase_          =							VersatileDiffusionTextToImagePipeline.from_pretrained('''shi-labs/versatile-diffusion''')
        # remove text_unet
        pipe.remove_unused_weights()
        pipe.to(_snake_case)
        pipe.set_progress_bar_config(disable=_snake_case)

        UpperCAmelCase_          =							'''A painting of a squirrel eating a burger '''
        UpperCAmelCase_          =							torch.manual_seed(0)
        UpperCAmelCase_          =							pipe(
            prompt=_snake_case					,						generator=_snake_case					,						guidance_scale=7.5					,						num_inference_steps=2					,						output_type='''numpy''').images

        with tempfile.TemporaryDirectory() as tmpdirname:
             pipe.save_pretrained(_snake_case)
             UpperCAmelCase_          =							VersatileDiffusionTextToImagePipeline.from_pretrained(_snake_case)
        pipe.to(_snake_case)
        pipe.set_progress_bar_config(disable=_snake_case)

        UpperCAmelCase_          =							generator.manual_seed(0)
        UpperCAmelCase_          =							pipe(
            prompt=_snake_case					,						generator=_snake_case					,						guidance_scale=7.5					,						num_inference_steps=2					,						output_type='''numpy''').images

        assert np.abs(image - new_image).sum() < 1e-5, "Models don't have the same forward pass"







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							VersatileDiffusionTextToImagePipeline.from_pretrained(
            '''shi-labs/versatile-diffusion'''					,						torch_dtype=torch.floataa)
        pipe.to(_snake_case)
        pipe.set_progress_bar_config(disable=_snake_case)

        UpperCAmelCase_          =							'''A painting of a squirrel eating a burger '''
        UpperCAmelCase_          =							torch.manual_seed(0)
        UpperCAmelCase_          =							pipe(
            prompt=_snake_case					,						generator=_snake_case					,						guidance_scale=7.5					,						num_inference_steps=50					,						output_type='''numpy''').images

        UpperCAmelCase_          =							image[0, 253:256, 253:256, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_          =							np.array([0.3_3_6_7, 0.3_1_6_9, 0.2_6_5_6, 0.3_8_7_0, 0.4_7_9_0, 0.3_7_9_6, 0.4_0_0_9, 0.4_8_7_8, 0.4_7_7_8])

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












snake_case_ : Dict    		=       {
    "Pillow": "Pillow",
    "accelerate": "accelerate>=0.11.0",
    "compel": "compel==0.1.8",
    "black": "black~=23.1",
    "datasets": "datasets",
    "filelock": "filelock",
    "flax": "flax>=0.4.1",
    "hf-doc-builder": "hf-doc-builder>=0.3.0",
    "huggingface-hub": "huggingface-hub>=0.13.2",
    "requests-mock": "requests-mock==1.10.0",
    "importlib_metadata": "importlib_metadata",
    "invisible-watermark": "invisible-watermark",
    "isort": "isort>=5.5.4",
    "jax": "jax>=0.2.8,!=0.3.2",
    "jaxlib": "jaxlib>=0.1.65",
    "Jinja2": "Jinja2",
    "k-diffusion": "k-diffusion>=0.0.12",
    "torchsde": "torchsde",
    "note_seq": "note_seq",
    "librosa": "librosa",
    "numpy": "numpy",
    "omegaconf": "omegaconf",
    "parameterized": "parameterized",
    "protobuf": "protobuf>=3.20.3,<4",
    "pytest": "pytest",
    "pytest-timeout": "pytest-timeout",
    "pytest-xdist": "pytest-xdist",
    "ruff": "ruff>=0.0.241",
    "safetensors": "safetensors",
    "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
    "scipy": "scipy",
    "onnx": "onnx",
    "regex": "regex!=2019.12.17",
    "requests": "requests",
    "tensorboard": "tensorboard",
    "torch": "torch>=1.4",
    "torchvision": "torchvision",
    "transformers": "transformers>=4.25.1",
    "urllib3": "urllib3<=2.0.0",
}






