






'''simple docstring'''


import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
@slow
@require_torch_gpu
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):






			def 				A      (   self							:						int			)							->    int:

									'''simple docstring'''







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






			def 				A      (   self							:						Any			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = StableDiffusionKDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""			)
									UpperCamelCase__         = sd_pipe.to(lowercase			)
									sd_pipe.set_progress_bar_config(disable=lowercase			)

									sd_pipe.set_scheduler("""sample_euler"""			)

									UpperCamelCase__         = """A painting of a squirrel eating a burger"""
									UpperCamelCase__         = torch.manual_seed(0			)
									UpperCamelCase__         = sd_pipe([prompt] ,      generator=lowercase ,      guidance_scale=9.0 ,      num_inference_steps=2_0 ,      output_type="""np"""			)

									UpperCamelCase__         = output.images

									UpperCamelCase__         = image[0, -3:, -3:, -1]

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									UpperCamelCase__         = np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]			)

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






			def 				A      (   self							:						Tuple			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = StableDiffusionKDiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""			)
									UpperCamelCase__         = sd_pipe.to(lowercase			)
									sd_pipe.set_progress_bar_config(disable=lowercase			)

									sd_pipe.set_scheduler("""sample_euler"""			)

									UpperCamelCase__         = """A painting of a squirrel eating a burger"""
									UpperCamelCase__         = torch.manual_seed(0			)
									UpperCamelCase__         = sd_pipe([prompt] ,      generator=lowercase ,      guidance_scale=9.0 ,      num_inference_steps=2_0 ,      output_type="""np"""			)

									UpperCamelCase__         = output.images

									UpperCamelCase__         = image[0, -3:, -3:, -1]

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									UpperCamelCase__         = np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]			)

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







			def 				A      (   self							:						Dict			)							->    Union[str, Any]:

									'''simple docstring'''







									UpperCamelCase__         = StableDiffusionKDiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""			)
									UpperCamelCase__         = sd_pipe.to(lowercase			)
									sd_pipe.set_progress_bar_config(disable=lowercase			)

									sd_pipe.set_scheduler("""sample_dpmpp_2m"""			)

									UpperCamelCase__         = """A painting of a squirrel eating a burger"""
									UpperCamelCase__         = torch.manual_seed(0			)
									UpperCamelCase__         = sd_pipe(
									    [prompt] ,      generator=lowercase ,      guidance_scale=7.5 ,      num_inference_steps=1_5 ,      output_type="""np""" ,      use_karras_sigmas=lowercase ,      )

									UpperCamelCase__         = output.images

									UpperCamelCase__         = image[0, -3:, -3:, -1]

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									UpperCamelCase__         = np.array(
									    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]			)

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








'''simple docstring'''


import numpy as np

from transformers import Pipeline







def 	__magic_name__( _A					):
						'''simple docstring'''





						UpperCamelCase__         = np.max(_A						,      axis=-1						,      keepdims=_A					)
						UpperCamelCase__         = np.exp(outputs - maxes					)
						return shifted_exp / shifted_exp.sum(axis=-1						,      keepdims=_A					)







class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			def 				A      (   self							:						int ,      **lowercase							:						List[Any]			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = {}
									if "second_text" in kwargs:
															UpperCamelCase__         = kwargs["""second_text"""]
									return preprocess_kwargs, {}, {}






			def 				A      (   self							:						Optional[int] ,      lowercase							:						Union[str, Any] ,      lowercase							:						Any=None			)							->    Optional[int]:

									'''simple docstring'''







									return self.tokenizer(lowercase ,      text_pair=lowercase ,      return_tensors=self.framework			)






			def 				A      (   self							:						Union[str, Any] ,      lowercase							:						Optional[int]			)							->    Any:

									'''simple docstring'''







									return self.model(**lowercase			)







			def 				A      (   self							:						List[str] ,      lowercase							:						Any			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = model_outputs.logits[0].numpy()
									UpperCamelCase__         = softmax(lowercase			)

									UpperCamelCase__         = np.argmax(lowercase			)
									UpperCamelCase__         = self.model.config.idalabel[best_class]
									UpperCamelCase__         = probabilities[best_class].item()
									UpperCamelCase__         = logits.tolist()
									return {"label": label, "score": score, "logits": logits}
