






'''simple docstring'''




import os
import sys
import warnings
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen
from ..table import array_cast
from ..utils.file_utils import is_local_path
from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
		import PIL.Image

		from .features import FeatureType


lowerCAmelCase__          =						None
lowerCAmelCase__          =						'<' if sys.byteorder == 'little' else '>'
# Origin: https://github.com/python-pillow/Pillow/blob/698951e19e19972aeed56df686868f1329981c12/src/PIL/Image.py#L3126 minus "|i1" which values are not preserved correctly when saving and loading an image
lowerCAmelCase__          =						[
    np.dtype('|b1'),
    np.dtype('|u1'),
    np.dtype('<u2'),
    np.dtype('>u2'),
    np.dtype('<i2'),
    np.dtype('>i2'),
    np.dtype('<u4'),
    np.dtype('>u4'),
    np.dtype('<i4'),
    np.dtype('>i4'),
    np.dtype('<f4'),
    np.dtype('>f4'),
    np.dtype('<f8'),
    np.dtype('>f8'),
]







@dataclass
class       __lowercase					:
	_lowerCamelCase              =  True
	_lowerCamelCase              =  None
	# Automatically constructed
	_lowerCamelCase              =  "PIL.Image.Image"
	_lowerCamelCase              =  pa.struct({'''bytes''': pa.binary(), '''path''': pa.string()}    )
	_lowerCamelCase              =  field(default='''Image'''   ,				init=__lowerCamelCase   ,				repr=__lowerCamelCase    )

	def __call__(      self   :						Union[str, Any]):
			return self.pa_type

	def      __UpperCamelCase							(      self   :						Any  ,					UpperCAmelCase_   :						Union[str, bytes, dict, np.ndarray, "PIL.Image.Image"]):
			if config.PIL_AVAILABLE:
					import PIL.Image
			else:
					raise ImportError('To support encoding images, please install \'Pillow\'.')

			if isinstance(UpperCAmelCase_  ,					UpperCAmelCase_):
					UpperCamelCase__				:					int					=					np.array(UpperCAmelCase_)

			if isinstance(UpperCAmelCase_  ,					UpperCAmelCase_):
					return {"path": value, "bytes": None}
			elif isinstance(UpperCAmelCase_  ,					UpperCAmelCase_):
					return {"path": None, "bytes": value}
			elif isinstance(UpperCAmelCase_  ,					np.ndarray):
					# convert the image array to PNG/TIFF bytes
					return encode_np_array(UpperCAmelCase_)
			elif isinstance(UpperCAmelCase_  ,					PIL.Image.Image):
					# convert the PIL image to bytes (default format is PNG/TIFF)
					return encode_pil_image(UpperCAmelCase_)
			elif value.get('path') is not None and os.path.isfile(value['path']):
					# we set "bytes": None to not duplicate the data if they're already available locally
					return {"bytes": None, "path": value.get('path')}
			elif value.get('bytes') is not None or value.get('path') is not None:
					# store the image bytes, and path is used to infer the image format using the file extension
					return {"bytes": value.get('bytes'), "path": value.get('path')}
			else:
					raise ValueError(
					    F'An image sample should have one of \'path\' or \'bytes\' but they are missing or None in {value}.')

	def      __UpperCamelCase							(      self   :						Optional[Any]  ,					UpperCAmelCase_   :						dict  ,					UpperCAmelCase_   :						int=None):
			if not self.decode:
					raise RuntimeError('Decoding is disabled for this feature. Please use Image(decode=True) instead.')

			if config.PIL_AVAILABLE:
					import PIL.Image
			else:
					raise ImportError('To support decoding images, please install \'Pillow\'.')

			if token_per_repo_id is None:
					UpperCamelCase__				:					int					=					{}

			UpperCamelCase__,	UpperCamelCase__				:					int					=					value['path'], value['bytes']
			if bytes_ is None:
					if path is None:
							raise ValueError(F'An image should have one of \'path\' or \'bytes\' but both are None in {value}.')
					else:
							if is_local_path(UpperCAmelCase_):
									UpperCamelCase__				:					int					=					PIL.Image.open(UpperCAmelCase_)
							else:
									UpperCamelCase__				:					str					=					path.split('::')[-1]
									try:
											UpperCamelCase__				:					Union[str, Any]					=					string_to_dict(UpperCAmelCase_  ,					config.HUB_DATASETS_URL)['repo_id']
											UpperCamelCase__				:					Optional[Any]					=					token_per_repo_id.get(UpperCAmelCase_)
									except ValueError:
											UpperCamelCase__				:					Any					=					None
									with xopen(UpperCAmelCase_  ,					'rb'  ,					use_auth_token=UpperCAmelCase_) as f:
											UpperCamelCase__				:					Optional[int]					=					BytesIO(f.read())
									UpperCamelCase__				:					Tuple					=					PIL.Image.open(bytes_)
			else:
					UpperCamelCase__				:					List[Any]					=					PIL.Image.open(BytesIO(bytes_))
			image.load()  # to avoid "Too many open files" errors
			return image

	def      __UpperCamelCase							(      self   :						List[str]):
			from .features import Value

			return (
			    self
			    if self.decode
			    else {
			        "bytes": Value('binary'),
			        "path": Value('string'),
			    }
			)

	def      __UpperCamelCase							(      self   :						Tuple  ,					UpperCAmelCase_   :						Union[pa.StringArray, pa.StructArray, pa.ListArray]):
			if pa.types.is_string(storage.type):
					UpperCamelCase__				:					Dict					=					pa.array([None] * len(UpperCAmelCase_)  ,					type=pa.binary())
					UpperCamelCase__				:					Union[str, Any]					=					pa.StructArray.from_arrays([bytes_array, storage]  ,					['bytes', 'path']  ,					mask=storage.is_null())
			elif pa.types.is_binary(storage.type):
					UpperCamelCase__				:					Tuple					=					pa.array([None] * len(UpperCAmelCase_)  ,					type=pa.string())
					UpperCamelCase__				:					Optional[int]					=					pa.StructArray.from_arrays([storage, path_array]  ,					['bytes', 'path']  ,					mask=storage.is_null())
			elif pa.types.is_struct(storage.type):
					if storage.type.get_field_index('bytes') >= 0:
							UpperCamelCase__				:					Tuple					=					storage.field('bytes')
					else:
							UpperCamelCase__				:					Tuple					=					pa.array([None] * len(UpperCAmelCase_)  ,					type=pa.binary())
					if storage.type.get_field_index('path') >= 0:
							UpperCamelCase__				:					int					=					storage.field('path')
					else:
							UpperCamelCase__				:					List[str]					=					pa.array([None] * len(UpperCAmelCase_)  ,					type=pa.string())
					UpperCamelCase__				:					Dict					=					pa.StructArray.from_arrays([bytes_array, path_array]  ,					['bytes', 'path']  ,					mask=storage.is_null())
			elif pa.types.is_list(storage.type):
					UpperCamelCase__				:					Optional[Any]					=					pa.array(
					    [encode_np_array(np.array(UpperCAmelCase_))['bytes'] if arr is not None else None for arr in storage.to_pylist()]  ,					type=pa.binary()  ,					)
					UpperCamelCase__				:					Optional[Any]					=					pa.array([None] * len(UpperCAmelCase_)  ,					type=pa.string())
					UpperCamelCase__				:					Optional[int]					=					pa.StructArray.from_arrays(
					    [bytes_array, path_array]  ,					['bytes', 'path']  ,					mask=bytes_array.is_null())
			return array_cast(UpperCAmelCase_  ,					self.pa_type)





	def      __UpperCamelCase							(      self   :						str  ,					UpperCAmelCase_   :						pa.StructArray):

			@no_op_if_value_is_null
			def path_to_bytes(UpperCAmelCase_   :						int):
					with xopen(UpperCAmelCase_  ,					'rb') as f:
							UpperCamelCase__				:					List[str]					=					f.read()
					return bytes_

			UpperCamelCase__				:					Optional[Any]					=					pa.array(
			    [
			        (path_to_bytes(x['path']) if x['bytes'] is None else x['bytes']) if x is not None else None
			        for x in storage.to_pylist()
			    ]  ,					type=pa.binary()  ,					)
			UpperCamelCase__				:					List[Any]					=					pa.array(
			    [os.path.basename(UpperCAmelCase_) if path is not None else None for path in storage.field('path').to_pylist()]  ,					type=pa.string()  ,					)
			UpperCamelCase__				:					str					=					pa.StructArray.from_arrays([bytes_array, path_array]  ,					['bytes', 'path']  ,					mask=bytes_array.is_null())
			return array_cast(UpperCAmelCase_  ,					self.pa_type)







def       __UpperCAmelCase						(  )							->     List[str]:
		if config.PIL_AVAILABLE:
				import PIL.Image
		else:
				raise ImportError('To support encoding images, please install \'Pillow\'.')

		global _IMAGE_COMPRESSION_FORMATS
		if _IMAGE_COMPRESSION_FORMATS is None:
				PIL.Image.init()
				UpperCamelCase__				:					Tuple					=					list(set(PIL.Image.OPEN.keys()) & set(PIL.Image.SAVE.keys()))
		return _IMAGE_COMPRESSION_FORMATS







def       __UpperCAmelCase						(  lowerCamelCase_)							->     bytes:
		UpperCamelCase__				:					List[Any]					=					BytesIO()
		if image.format in list_image_compression_formats():
				UpperCamelCase__				:					Dict					=					image.format
		else:
				UpperCamelCase__				:					List[str]					=					'PNG' if image.mode in ['1', 'L', 'LA', 'RGB', 'RGBA'] else 'TIFF'
		image.save(lowerCamelCase_     ,      format=lowerCamelCase_)
		return buffer.getvalue()







def       __UpperCAmelCase						(  lowerCamelCase_)							->     dict:
		if hasattr(lowerCamelCase_     ,      'filename') and image.filename != "":
				return {"path": image.filename, "bytes": None}
		else:
				return {"path": None, "bytes": image_to_bytes(lowerCamelCase_)}







def       __UpperCAmelCase						(  lowerCamelCase_)							->     dict:
		if config.PIL_AVAILABLE:
				import PIL.Image
		else:
				raise ImportError('To support encoding images, please install \'Pillow\'.')

		UpperCamelCase__				:					List[Any]					=					array.dtype
		UpperCamelCase__				:					Union[str, Any]					=					dtype.byteorder if dtype.byteorder != '=' else _NATIVE_BYTEORDER
		UpperCamelCase__				:					Optional[int]					=					dtype.kind
		UpperCamelCase__				:					int					=					dtype.itemsize

		UpperCamelCase__				:					List[str]					=					None

		# Multi-channel array case (only np.dtype("|u1") is allowed)
		if array.shape[2:]:
				UpperCamelCase__				:					Union[str, Any]					=					np.dtype('|u1')
				if dtype_kind not in ["u", "i"]:
						raise TypeError(
						    f'Unsupported array dtype {dtype} for image encoding. Only {dest_dtype} is supported for multi-channel arrays.')
				if dtype is not dest_dtype:
						warnings.warn(f'Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\'')
    # Exact match
		elif dtype in _VALID_IMAGE_ARRAY_DTPYES:
				UpperCamelCase__				:					Union[str, Any]					=					dtype
		else:  # Downcast the type within the kind (np.can_cast(from_type, to_type, casting="same_kind") doesn't behave as expected, so do it manually)
				while dtype_itemsize >= 1:
						UpperCamelCase__				:					List[str]					=					dtype_byteorder + dtype_kind + str(lowerCamelCase_)
						UpperCamelCase__				:					Any					=					np.dtype(lowerCamelCase_)
						if dest_dtype in _VALID_IMAGE_ARRAY_DTPYES:
								warnings.warn(f'Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\'')
								break
						else:
								dtype_itemsize //= 2

		if dest_dtype is None:
				raise TypeError(
				    f'Cannot convert dtype {dtype} to a valid image dtype. Valid image dtypes: {_VALID_IMAGE_ARRAY_DTPYES}')
		UpperCamelCase__				:					Optional[int]					=					PIL.Image.fromarray(array.astype(lowerCamelCase_))
		return {"path": None, "bytes": image_to_bytes(lowerCamelCase_)}







def       __UpperCAmelCase						(  lowerCamelCase_)							->     List[dict]:
		if config.PIL_AVAILABLE:
				import PIL.Image
		else:
				raise ImportError('To support encoding images, please install \'Pillow\'.')

		if objs:
				UpperCamelCase__,	UpperCamelCase__				:					str					=					first_non_null_value(lowerCamelCase_)
				if isinstance(lowerCamelCase_     ,      lowerCamelCase_):
						return [{"path": obj, "bytes": None} if obj is not None else None for obj in objs]
				if isinstance(lowerCamelCase_     ,      np.ndarray):
						UpperCamelCase__				:					List[str]					=					no_op_if_value_is_null(lowerCamelCase_)
						return [obj_to_image_dict_func(lowerCamelCase_) for obj in objs]
				elif isinstance(lowerCamelCase_     ,      PIL.Image.Image):
						UpperCamelCase__				:					Union[str, Any]					=					no_op_if_value_is_null(lowerCamelCase_)
						return [obj_to_image_dict_func(lowerCamelCase_) for obj in objs]
				else:
						return objs
		else:
				return objs












'''simple docstring'''




import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu







class       __lowercase					(unittest.TestCase    ):

	def      __UpperCamelCase							(      self   :						List[str]):
			# clean up the VRAM after each test
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()

	@property
	def      __UpperCamelCase							(      self   :						List[Any]):
			UpperCamelCase__				:					Union[str, Any]					=					1
			UpperCamelCase__				:					Union[str, Any]					=					3
			UpperCamelCase__				:					Dict					=					(32, 32)

			UpperCamelCase__				:					int					=					floats_tensor((batch_size, num_channels) + sizes  ,					rng=random.Random(0)).to(UpperCAmelCase_)
			return image

	@property
	def      __UpperCamelCase							(      self   :						Any):
			torch.manual_seed(0)
			UpperCamelCase__				:					Any					=					UNetaDConditionModel(
			    block_out_channels=(32, 64)  ,					layers_per_block=2  ,					sample_size=32  ,					in_channels=4  ,					out_channels=4  ,					down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,					up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,					cross_attention_dim=32  ,					)
			return model

	@property
	def      __UpperCamelCase							(      self   :						Any):
			torch.manual_seed(0)
			UpperCamelCase__				:					List[str]					=					AutoencoderKL(
			    block_out_channels=[32, 64]  ,					in_channels=3  ,					out_channels=3  ,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,					latent_channels=4  ,					)
			return model

	@property
	def      __UpperCamelCase							(      self   :						str):
			torch.manual_seed(0)
			UpperCamelCase__				:					Tuple					=					CLIPTextConfig(
			    bos_token_id=0  ,					eos_token_id=2  ,					hidden_size=32  ,					intermediate_size=37  ,					layer_norm_eps=1e-05  ,					num_attention_heads=4  ,					num_hidden_layers=5  ,					pad_token_id=1  ,					vocab_size=1_000  ,					)
			return CLIPTextModel(UpperCAmelCase_)

	@property
	def      __UpperCamelCase							(      self   :						Optional[Any]):
			def extract(*UpperCAmelCase_   :						List[Any]  ,					**UpperCAmelCase_   :						Dict):







					class       __lowercase					:

						def __init__(      self   :						List[Any]):
								UpperCamelCase__				:					Optional[Any]					=					torch.ones([0])





						def      __UpperCamelCase							(      self   :						Dict  ,					UpperCAmelCase_   :						int):
								self.pixel_values.to(UpperCAmelCase_)
								return self


					return Out()

			return extract

	def      __UpperCamelCase							(      self   :						str):
			UpperCamelCase__				:					Optional[Any]					=					'cpu'  # ensure determinism for the device-dependent torch.Generator
			UpperCamelCase__				:					Any					=					self.dummy_cond_unet
			UpperCamelCase__				:					Any					=					DDIMScheduler(
			    beta_start=0.0_00_85  ,					beta_end=0.0_12  ,					beta_schedule='scaled_linear'  ,					clip_sample=UpperCAmelCase_  ,					set_alpha_to_one=UpperCAmelCase_  ,					)

			UpperCamelCase__				:					List[str]					=					self.dummy_vae
			UpperCamelCase__				:					str					=					self.dummy_text_encoder
			UpperCamelCase__				:					Tuple					=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

			# make sure here that pndm scheduler skips prk
			UpperCamelCase__				:					Optional[Any]					=					StableDiffusionPipeline(
			    unet=UpperCAmelCase_  ,					scheduler=UpperCAmelCase_  ,					vae=UpperCAmelCase_  ,					text_encoder=UpperCAmelCase_  ,					tokenizer=UpperCAmelCase_  ,					safety_checker=UpperCAmelCase_  ,					feature_extractor=self.dummy_extractor  ,					)
			UpperCamelCase__				:					Optional[Any]					=					sd_pipe.to(UpperCAmelCase_)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					Optional[Any]					=					'A painting of a squirrel eating a burger'

			UpperCamelCase__				:					Dict					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					List[Any]					=					sd_pipe([prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=6.0  ,					num_inference_steps=2  ,					output_type='np')
			UpperCamelCase__				:					Tuple					=					output.images

			UpperCamelCase__				:					List[Any]					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					Tuple					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=6.0  ,					num_inference_steps=2  ,					output_type='np'  ,					return_dict=UpperCAmelCase_  ,					)[0]

			UpperCamelCase__				:					List[str]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Optional[int]					=					image_from_tuple[0, -3:, -3:, -1]

			assert image.shape == (1, 64, 64, 3)
			UpperCamelCase__				:					List[Any]					=					np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64])

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

	def      __UpperCamelCase							(      self   :						Dict):
			UpperCamelCase__				:					List[Any]					=					'cpu'  # ensure determinism for the device-dependent torch.Generator
			UpperCamelCase__				:					int					=					self.dummy_cond_unet
			UpperCamelCase__				:					Dict					=					PNDMScheduler(skip_prk_steps=UpperCAmelCase_)
			UpperCamelCase__				:					Optional[int]					=					self.dummy_vae
			UpperCamelCase__				:					Optional[int]					=					self.dummy_text_encoder
			UpperCamelCase__				:					Union[str, Any]					=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

			# make sure here that pndm scheduler skips prk
			UpperCamelCase__				:					Dict					=					StableDiffusionPipeline(
			    unet=UpperCAmelCase_  ,					scheduler=UpperCAmelCase_  ,					vae=UpperCAmelCase_  ,					text_encoder=UpperCAmelCase_  ,					tokenizer=UpperCAmelCase_  ,					safety_checker=UpperCAmelCase_  ,					feature_extractor=self.dummy_extractor  ,					)
			UpperCamelCase__				:					Tuple					=					sd_pipe.to(UpperCAmelCase_)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					List[str]					=					'A painting of a squirrel eating a burger'
			UpperCamelCase__				:					Union[str, Any]					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					str					=					sd_pipe([prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=6.0  ,					num_inference_steps=2  ,					output_type='np')

			UpperCamelCase__				:					List[str]					=					output.images

			UpperCamelCase__				:					Any					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					Optional[Any]					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=6.0  ,					num_inference_steps=2  ,					output_type='np'  ,					return_dict=UpperCAmelCase_  ,					)[0]

			UpperCamelCase__				:					Tuple					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Optional[Any]					=					image_from_tuple[0, -3:, -3:, -1]

			assert image.shape == (1, 64, 64, 3)
			UpperCamelCase__				:					List[Any]					=					np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93])

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

	def      __UpperCamelCase							(      self   :						Dict):
			UpperCamelCase__				:					Dict					=					StableDiffusionPipeline.from_pretrained(
			    'hf-internal-testing/tiny-stable-diffusion-lms-pipe'  ,					safety_checker=UpperCAmelCase_)
			assert isinstance(UpperCAmelCase_  ,					UpperCAmelCase_)
			assert isinstance(pipe.scheduler  ,					UpperCAmelCase_)
			assert pipe.safety_checker is None

			UpperCamelCase__				:					List[Any]					=					pipe('example prompt'  ,					num_inference_steps=2).images[0]
			assert image is not None

			# check that there's no error when saving a pipeline with one of the models being None
			with tempfile.TemporaryDirectory() as tmpdirname:
					pipe.save_pretrained(UpperCAmelCase_)
					UpperCamelCase__				:					List[str]					=					StableDiffusionPipeline.from_pretrained(UpperCAmelCase_)

			# sanity check that the pipeline still works
			assert pipe.safety_checker is None
			UpperCamelCase__				:					Optional[Any]					=					pipe('example prompt'  ,					num_inference_steps=2).images[0]
			assert image is not None





	@unittest.skipIf(torch_device != 'cuda'  ,					'This test requires a GPU')
	def      __UpperCamelCase							(      self   :						List[Any]):
			UpperCamelCase__				:					Dict					=					self.dummy_cond_unet
			UpperCamelCase__				:					str					=					PNDMScheduler(skip_prk_steps=UpperCAmelCase_)
			UpperCamelCase__				:					Any					=					self.dummy_vae
			UpperCamelCase__				:					Optional[Any]					=					self.dummy_text_encoder
			UpperCamelCase__				:					str					=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

			# put models in fp16
			UpperCamelCase__				:					Any					=					unet.half()
			UpperCamelCase__				:					Tuple					=					vae.half()
			UpperCamelCase__				:					Optional[int]					=					bert.half()

			# make sure here that pndm scheduler skips prk
			UpperCamelCase__				:					Optional[int]					=					StableDiffusionPipeline(
			    unet=UpperCAmelCase_  ,					scheduler=UpperCAmelCase_  ,					vae=UpperCAmelCase_  ,					text_encoder=UpperCAmelCase_  ,					tokenizer=UpperCAmelCase_  ,					safety_checker=UpperCAmelCase_  ,					feature_extractor=self.dummy_extractor  ,					)
			UpperCamelCase__				:					Dict					=					sd_pipe.to(UpperCAmelCase_)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					Any					=					'A painting of a squirrel eating a burger'
			UpperCamelCase__				:					int					=					sd_pipe([prompt]  ,					num_inference_steps=2  ,					output_type='np').images

			assert image.shape == (1, 64, 64, 3)







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

	def      __UpperCamelCase							(      self   :						Optional[int]):
			# clean up the VRAM after each test
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()

	def      __UpperCamelCase							(      self   :						List[Any]):
			UpperCamelCase__				:					Optional[int]					=					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'  ,					safety_checker=UpperCAmelCase_)
			UpperCamelCase__				:					Union[str, Any]					=					LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config)
			UpperCamelCase__				:					Optional[Any]					=					sd_pipe.to(UpperCAmelCase_)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					List[Any]					=					(
			    'portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle'
			    ' coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with'
			    ' anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and'
			    ' children from bahnhof zoo, detailed '
			)
			UpperCamelCase__				:					Any					=					4_003_660_346
			UpperCamelCase__				:					Any					=					7

			# without safety guidance (sld_guidance_scale = 0)
			UpperCamelCase__				:					int					=					torch.manual_seed(UpperCAmelCase_)
			UpperCamelCase__				:					Optional[int]					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=UpperCAmelCase_  ,					num_inference_steps=50  ,					output_type='np'  ,					width=512  ,					height=512  ,					sld_guidance_scale=0  ,					)

			UpperCamelCase__				:					str					=					output.images
			UpperCamelCase__				:					Union[str, Any]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Tuple					=					[0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

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

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

			# without safety guidance (strong configuration)
			UpperCamelCase__				:					Tuple					=					torch.manual_seed(UpperCAmelCase_)
			UpperCamelCase__				:					str					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=UpperCAmelCase_  ,					num_inference_steps=50  ,					output_type='np'  ,					width=512  ,					height=512  ,					sld_guidance_scale=2_000  ,					sld_warmup_steps=7  ,					sld_threshold=0.0_25  ,					sld_momentum_scale=0.5  ,					sld_mom_beta=0.7  ,					)

			UpperCamelCase__				:					Dict					=					output.images
			UpperCamelCase__				:					str					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Tuple					=					[0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

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

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

	def      __UpperCamelCase							(      self   :						Optional[Any]):
			UpperCamelCase__				:					Dict					=					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'  ,					safety_checker=UpperCAmelCase_)
			UpperCamelCase__				:					str					=					LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config)
			UpperCamelCase__				:					Dict					=					sd_pipe.to(UpperCAmelCase_)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					str					=					'padme amidala taking a bath artwork, safe for work, no nudity'
			UpperCamelCase__				:					Tuple					=					2_734_971_755
			UpperCamelCase__				:					Tuple					=					7

			UpperCamelCase__				:					Tuple					=					torch.manual_seed(UpperCAmelCase_)
			UpperCamelCase__				:					int					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=UpperCAmelCase_  ,					num_inference_steps=50  ,					output_type='np'  ,					width=512  ,					height=512  ,					sld_guidance_scale=0  ,					)

			UpperCamelCase__				:					int					=					output.images
			UpperCamelCase__				:					Union[str, Any]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Any					=					[0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

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

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

			UpperCamelCase__				:					List[str]					=					torch.manual_seed(UpperCAmelCase_)
			UpperCamelCase__				:					Union[str, Any]					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=UpperCAmelCase_  ,					num_inference_steps=50  ,					output_type='np'  ,					width=512  ,					height=512  ,					sld_guidance_scale=2_000  ,					sld_warmup_steps=7  ,					sld_threshold=0.0_25  ,					sld_momentum_scale=0.5  ,					sld_mom_beta=0.7  ,					)

			UpperCamelCase__				:					Tuple					=					output.images
			UpperCamelCase__				:					List[str]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Union[str, Any]					=					[0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

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

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





	def      __UpperCamelCase							(      self   :						Any):
			UpperCamelCase__				:					Optional[Any]					=					StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5')
			UpperCamelCase__				:					Optional[Any]					=					sd_pipe.to(UpperCAmelCase_)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					int					=					(
			    'the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c.'
			    ' leyendecker'
			)
			UpperCamelCase__				:					Any					=					1_044_355_234
			UpperCamelCase__				:					Optional[int]					=					12

			UpperCamelCase__				:					Optional[int]					=					torch.manual_seed(UpperCAmelCase_)
			UpperCamelCase__				:					str					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=UpperCAmelCase_  ,					num_inference_steps=50  ,					output_type='np'  ,					width=512  ,					height=512  ,					sld_guidance_scale=0  ,					)

			UpperCamelCase__				:					List[str]					=					output.images
			UpperCamelCase__				:					Any					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					str					=					np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])

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

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

			UpperCamelCase__				:					int					=					torch.manual_seed(UpperCAmelCase_)
			UpperCamelCase__				:					List[str]					=					sd_pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					guidance_scale=UpperCAmelCase_  ,					num_inference_steps=50  ,					output_type='np'  ,					width=512  ,					height=512  ,					sld_guidance_scale=2_000  ,					sld_warmup_steps=7  ,					sld_threshold=0.0_25  ,					sld_momentum_scale=0.5  ,					sld_mom_beta=0.7  ,					)

			UpperCamelCase__				:					Optional[Any]					=					output.images
			UpperCamelCase__				:					List[Any]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					str					=					np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61])
			assert image.shape == (1, 512, 512, 3)

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




