


'''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.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class 						lowerCAmelCase__      (      unittest.TestCase						):




			def 				lowerCAmelCase__     (		self	:  List[Any]			) ->Optional[int]:



					'''simple docstring'''







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




			@property
			def 				lowerCAmelCase__     (		self	:  Dict			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							Any             =   1
					_UpperCAmelCase		:							Optional[int]             =   3
					_UpperCAmelCase		:							int             =   (32, 32)

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




			@property
			def 				lowerCAmelCase__     (		self	:  Dict			) ->List[str]:



					'''simple docstring'''







					torch.manual_seed(0			)
					_UpperCAmelCase		:							int             =   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 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Any:



					'''simple docstring'''







					torch.manual_seed(0			)
					_UpperCAmelCase		:							Any             =   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 				lowerCAmelCase__     (		self	:  Dict			) ->Union[str, Any]:



					'''simple docstring'''







					torch.manual_seed(0			)
					_UpperCAmelCase		:							List[Any]             =   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=10_00						,	)
					return CLIPTextModel(lowerCamelCase__			)




			@property
			def 				lowerCAmelCase__     (		self	:  int			) ->Optional[Any]:



					'''simple docstring'''







					def extract(*lowerCamelCase__	:  List[str]						,	**lowerCamelCase__	:  List[Any]			):






							class 						lowerCAmelCase__      :




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



												'''simple docstring'''







												_UpperCAmelCase		:							Dict             =   torch.ones([0]			)






										def 				lowerCAmelCase__     (		self	:  List[str]						,	lowerCamelCase__	:  int			) ->Any:



												'''simple docstring'''







												self.pixel_values.to(lowerCamelCase__			)
												return self




							return Out()

					return extract




			def 				lowerCAmelCase__     (		self	:  List[Any]			) ->int:



					'''simple docstring'''







					_UpperCAmelCase		:							str             =   "cpu"  # ensure determinism for the device-dependent torch.Generator
					_UpperCAmelCase		:							List[str]             =   self.dummy_cond_unet
					_UpperCAmelCase		:							int             =   DDIMScheduler(
					    beta_start=0.0_0_0_8_5						,	beta_end=0.0_1_2						,	beta_schedule="scaled_linear"						,	clip_sample=lowerCamelCase__						,	set_alpha_to_one=lowerCamelCase__						,	)

					_UpperCAmelCase		:							Tuple             =   self.dummy_vae
					_UpperCAmelCase		:							Any             =   self.dummy_text_encoder
					_UpperCAmelCase		:							Optional[Any]             =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

					# make sure here that pndm scheduler skips prk
					_UpperCAmelCase		:							List[Any]             =   StableDiffusionPipeline(
					    unet=lowerCamelCase__						,	scheduler=lowerCamelCase__						,	vae=lowerCamelCase__						,	text_encoder=lowerCamelCase__						,	tokenizer=lowerCamelCase__						,	safety_checker=lowerCamelCase__						,	feature_extractor=self.dummy_extractor						,	)
					_UpperCAmelCase		:							Union[str, Any]             =   sd_pipe.to(lowerCamelCase__			)
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase		:							str             =   "A painting of a squirrel eating a burger"

					_UpperCAmelCase		:							Optional[int]             =   torch.Generator(device=lowerCamelCase__			).manual_seed(0			)
					_UpperCAmelCase		:							Any             =   sd_pipe([prompt]						,	generator=lowerCamelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="np"			)
					_UpperCAmelCase		:							Any             =   output.images

					_UpperCAmelCase		:							int             =   torch.Generator(device=lowerCamelCase__			).manual_seed(0			)
					_UpperCAmelCase		:							Any             =   sd_pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="np"						,	return_dict=lowerCamelCase__						,	)[0]

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

					assert image.shape == (1, 64, 64, 3)
					_UpperCAmelCase		:							Dict             =   np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7]			)

					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 				lowerCAmelCase__     (		self	:  Optional[int]			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   "cpu"  # ensure determinism for the device-dependent torch.Generator
					_UpperCAmelCase		:							Dict             =   self.dummy_cond_unet
					_UpperCAmelCase		:							Any             =   PNDMScheduler(skip_prk_steps=lowerCamelCase__			)
					_UpperCAmelCase		:							List[Any]             =   self.dummy_vae
					_UpperCAmelCase		:							str             =   self.dummy_text_encoder
					_UpperCAmelCase		:							Optional[Any]             =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

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

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

					_UpperCAmelCase		:							Union[str, Any]             =   output.images

					_UpperCAmelCase		:							Union[str, Any]             =   torch.Generator(device=lowerCamelCase__			).manual_seed(0			)
					_UpperCAmelCase		:							Any             =   sd_pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="np"						,	return_dict=lowerCamelCase__						,	)[0]

					_UpperCAmelCase		:							List[str]             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							Tuple             =   image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)
					_UpperCAmelCase		:							Any             =   np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4]			)

					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 				lowerCAmelCase__     (		self	:  int			) ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   StableDiffusionPipeline.from_pretrained(
					    "hf-internal-testing/tiny-stable-diffusion-lms-pipe"						,	safety_checker=lowerCamelCase__			)
					assert isinstance(lowerCamelCase__						,	lowerCamelCase__			)
					assert isinstance(pipe.scheduler						,	lowerCamelCase__			)
					assert pipe.safety_checker is None

					_UpperCAmelCase		:							Optional[int]             =   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(lowerCamelCase__			)
							_UpperCAmelCase		:							Dict             =   StableDiffusionPipeline.from_pretrained(lowerCamelCase__			)

					# sanity check that the pipeline still works
					assert pipe.safety_checker is None
					_UpperCAmelCase		:							List[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 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   self.dummy_cond_unet
					_UpperCAmelCase		:							Optional[int]             =   PNDMScheduler(skip_prk_steps=lowerCamelCase__			)
					_UpperCAmelCase		:							str             =   self.dummy_vae
					_UpperCAmelCase		:							Optional[Any]             =   self.dummy_text_encoder
					_UpperCAmelCase		:							List[Any]             =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

					# put models in fp16
					_UpperCAmelCase		:							Union[str, Any]             =   unet.half()
					_UpperCAmelCase		:							str             =   vae.half()
					_UpperCAmelCase		:							List[str]             =   bert.half()

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

					_UpperCAmelCase		:							List[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 						lowerCAmelCase__      (      unittest.TestCase						):




			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->int:



					'''simple docstring'''







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




			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   "cuda"
					_UpperCAmelCase		:							List[Any]             =   StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			)
					_UpperCAmelCase		:							Dict             =   pipe.to(lowerCamelCase__			)
					pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase		:							Optional[int]             =   "a photo of a cat"
					_UpperCAmelCase		:							List[Any]             =   {
					    "editing_prompt": ["sunglasses"],
					    "reverse_editing_direction": [False],
					    "edit_warmup_steps": 10,
					    "edit_guidance_scale": 6,
					    "edit_threshold": 0.9_5,
					    "edit_momentum_scale": 0.5,
					    "edit_mom_beta": 0.6,
					}

					_UpperCAmelCase		:							Union[str, Any]             =   3
					_UpperCAmelCase		:							Optional[int]             =   7

					# no sega enabled
					_UpperCAmelCase		:							str             =   torch.Generator(lowerCamelCase__			)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							List[Any]             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	)

					_UpperCAmelCase		:							Optional[int]             =   output.images
					_UpperCAmelCase		:							Dict             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							Optional[Any]             =   [
					    0.3_4_6_7_3_1_1_3,
					    0.3_8_4_9_2_7_3_3,
					    0.3_7_5_9_7_3_5_2,
					    0.3_4_0_8_6_3_3_5,
					    0.3_5_6_5_0_7_4_8,
					    0.3_5_5_7_9_2_0_5,
					    0.3_3_8_4_7_6_3,
					    0.3_4_3_4_0_2_3_6,
					    0.3_5_7_3_2_7_1,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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

					# with sega enabled
					# generator = torch.manual_seed(seed)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							int             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	**lowerCamelCase__						,	)

					_UpperCAmelCase		:							Dict             =   output.images
					_UpperCAmelCase		:							Union[str, Any]             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							Optional[int]             =   [
					    0.4_1_8_8_7_8_2_6,
					    0.3_7_7_2_8_7_6_6,
					    0.3_0_1_3_8_2_7_2,
					    0.4_1_4_1_6_3_3_5,
					    0.4_1_6_6_4_9_8_5,
					    0.3_6_2_8_3_3_9_2,
					    0.3_6_1_9_1_2_4_6,
					    0.4_3_3_6_4_4_6_5,
					    0.4_3_0_0_1_7_3_2,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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




			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:							Tuple             =   "cuda"
					_UpperCAmelCase		:							Any             =   StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			)
					_UpperCAmelCase		:							List[Any]             =   pipe.to(lowerCamelCase__			)
					pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase		:							Tuple             =   "an image of a crowded boulevard, realistic, 4k"
					_UpperCAmelCase		:							int             =   {
					    "editing_prompt": "crowd, crowded, people",
					    "reverse_editing_direction": True,
					    "edit_warmup_steps": 10,
					    "edit_guidance_scale": 8.3,
					    "edit_threshold": 0.9,
					    "edit_momentum_scale": 0.5,
					    "edit_mom_beta": 0.6,
					}

					_UpperCAmelCase		:							List[Any]             =   9
					_UpperCAmelCase		:							Tuple             =   7

					# no sega enabled
					_UpperCAmelCase		:							Tuple             =   torch.Generator(lowerCamelCase__			)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							str             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	)

					_UpperCAmelCase		:							List[str]             =   output.images
					_UpperCAmelCase		:							Union[str, Any]             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							str             =   [
					    0.4_3_4_9_7_9_9_8,
					    0.9_1_8_1_4_0_6_5,
					    0.7_5_4_0_7_3_9,
					    0.5_5_5_8_0_2_0_5,
					    0.8_4_6_7_2_6_5,
					    0.5_3_8_9_6_9_1,
					    0.6_2_5_7_4_5_0_6,
					    0.5_8_8_9_7_7_6_3,
					    0.5_0_9_2_6_7_5_7,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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

					# with sega enabled
					# generator = torch.manual_seed(seed)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							List[str]             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	**lowerCamelCase__						,	)

					_UpperCAmelCase		:							Union[str, Any]             =   output.images
					_UpperCAmelCase		:							Dict             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							Dict             =   [
					    0.3_0_8_9_7_1_9,
					    0.3_0_5_0_0_1_4_4,
					    0.2_9_0_1_6_0_4_2,
					    0.3_0_6_3_0_9_6_4,
					    0.3_2_5_6_8_7,
					    0.2_9_4_1_9_2_2_5,
					    0.2_9_0_8_0_9_1,
					    0.2_8_7_2_3_5_9_8,
					    0.2_7_6_9_6_2_9_4,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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




			def 				lowerCAmelCase__     (		self	:  List[str]			) ->int:



					'''simple docstring'''







					_UpperCAmelCase		:							Union[str, Any]             =   "cuda"
					_UpperCAmelCase		:							str             =   StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			)
					_UpperCAmelCase		:							Dict             =   pipe.to(lowerCamelCase__			)
					pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase		:							str             =   "a castle next to a river"
					_UpperCAmelCase		:							Optional[int]             =   {
					    "editing_prompt": ["boat on a river, boat", "monet, impression, sunrise"],
					    "reverse_editing_direction": False,
					    "edit_warmup_steps": [15, 18],
					    "edit_guidance_scale": 6,
					    "edit_threshold": [0.9, 0.8],
					    "edit_momentum_scale": 0.5,
					    "edit_mom_beta": 0.6,
					}

					_UpperCAmelCase		:							Tuple             =   48
					_UpperCAmelCase		:							List[Any]             =   7

					# no sega enabled
					_UpperCAmelCase		:							List[Any]             =   torch.Generator(lowerCamelCase__			)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							List[str]             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	)

					_UpperCAmelCase		:							List[str]             =   output.images
					_UpperCAmelCase		:							Tuple             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							str             =   [
					    0.7_5_1_6_3_5_5_5,
					    0.7_6_0_3_7_1_4_5,
					    0.6_1_7_8_5,
					    0.9_1_8_9_6_7_3,
					    0.8_6_2_7_7_0_1,
					    0.8_5_1_8_9_6_9_4,
					    0.8_5_1_2_8_1_3,
					    0.8_7_0_1_2_0_7_6,
					    0.8_3_1_2_8_5_7,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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

					# with sega enabled
					# generator = torch.manual_seed(seed)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							List[Any]             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	**lowerCamelCase__						,	)

					_UpperCAmelCase		:							Tuple             =   output.images
					_UpperCAmelCase		:							Optional[Any]             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							int             =   [
					    0.7_3_5_5_3_3_6_5,
					    0.7_5_3_7_2_7_1,
					    0.7_4_3_4_1_9_0_5,
					    0.6_6_4_8_0_3_5_6,
					    0.6_4_7_2_9_2_5,
					    0.6_3_0_3_9_4_1_6,
					    0.6_4_8_1_2_9_0_5,
					    0.6_7_4_9_7_1_7,
					    0.6_5_1_7_1_0_2,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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






			def 				lowerCAmelCase__     (		self	:  List[Any]			) ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							Tuple             =   "cuda"
					_UpperCAmelCase		:							Tuple             =   StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"						,	torch_dtype=torch.floataa			)
					_UpperCAmelCase		:							int             =   pipe.to(lowerCamelCase__			)
					pipe.set_progress_bar_config(disable=lowerCamelCase__			)

					_UpperCAmelCase		:							Optional[Any]             =   "a photo of a cat"
					_UpperCAmelCase		:							Any             =   {
					    "editing_prompt": ["sunglasses"],
					    "reverse_editing_direction": [False],
					    "edit_warmup_steps": 10,
					    "edit_guidance_scale": 6,
					    "edit_threshold": 0.9_5,
					    "edit_momentum_scale": 0.5,
					    "edit_mom_beta": 0.6,
					}

					_UpperCAmelCase		:							int             =   3
					_UpperCAmelCase		:							str             =   7

					# no sega enabled
					_UpperCAmelCase		:							Dict             =   torch.Generator(lowerCamelCase__			)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							Optional[int]             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	)

					_UpperCAmelCase		:							List[str]             =   output.images
					_UpperCAmelCase		:							str             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							Dict             =   [
					    0.3_4_8_8_7_6_9_5,
					    0.3_8_7_6_9_5_3,
					    0.3_7_5,
					    0.3_4_4_2_3_8_2_8,
					    0.3_5_8_1_5_4_3,
					    0.3_5_7_1_7_7_7_3,
					    0.3_3_8_3_7_8_9,
					    0.3_4_5_7_0_3_1_2,
					    0.3_5_9_3_7_5,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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

					# with sega enabled
					# generator = torch.manual_seed(seed)
					generator.manual_seed(lowerCamelCase__			)
					_UpperCAmelCase		:							Union[str, Any]             =   pipe(
					    [prompt]						,	generator=lowerCamelCase__						,	guidance_scale=lowerCamelCase__						,	num_inference_steps=50						,	output_type="np"						,	width=5_12						,	height=5_12						,	**lowerCamelCase__						,	)

					_UpperCAmelCase		:							Any             =   output.images
					_UpperCAmelCase		:							Any             =   image[0, -3:, -3:, -1]
					_UpperCAmelCase		:							Union[str, Any]             =   [
					    0.4_2_2_8_5_1_5_6,
					    0.3_6_9_1_4_0_6_2,
					    0.2_9_0_7_7_1_4_8,
					    0.4_2_0_4_1_0_1_6,
					    0.4_1_9_1_8_9_4_5,
					    0.3_5_4_9_8_0_4_7,
					    0.3_6_1_8_1_6_4,
					    0.4_4_2_3_8_2_8,
					    0.4_3_1_1_5_2_3_4,
					]

					assert image.shape == (1, 5_12, 5_12, 3)

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







'''simple docstring'''


import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin






class 						lowerCAmelCase__      (      UpperCAmelCase__     ,							UpperCAmelCase__     ,							unittest.TestCase						):
			lowerCAmelCase			:						Any							 =  UNetaDModel
			lowerCAmelCase			:						List[Any]							 =  "sample"




			@property
			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   4
					_UpperCAmelCase		:							Any             =   14
					_UpperCAmelCase		:							Any             =   16

					_UpperCAmelCase		:							Optional[int]             =   floats_tensor((batch_size, num_features, seq_len)			).to(lowerCamelCase__			)
					_UpperCAmelCase		:							Dict             =   torch.tensor([10] * batch_size			).to(lowerCamelCase__			)

					return {"sample": noise, "timestep": time_step}




			@property
			def 				lowerCAmelCase__     (		self	:  Tuple			) ->List[Any]:



					'''simple docstring'''







					return (4, 14, 16)




			@property
			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->Tuple:



					'''simple docstring'''







					return (4, 14, 16)




			def 				lowerCAmelCase__     (		self	:  int			) ->Tuple:



					'''simple docstring'''







					pass




			def 				lowerCAmelCase__     (		self	:  Dict			) ->Tuple:



					'''simple docstring'''







					pass




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  Any			) ->List[str]:



					'''simple docstring'''







					super().test_determinism()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->Dict:



					'''simple docstring'''







					super().test_outputs_equivalence()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  int			) ->List[str]:



					'''simple docstring'''







					super().test_from_save_pretrained()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Optional[int]:



					'''simple docstring'''







					super().test_from_save_pretrained_variant()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  str			) ->int:



					'''simple docstring'''







					super().test_model_from_pretrained()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  int			) ->Union[str, Any]:



					'''simple docstring'''







					super().test_output()




			def 				lowerCAmelCase__     (		self	:  int			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							Any             =   {
					    "block_out_channels": (32, 64, 1_28, 2_56),
					    "in_channels": 14,
					    "out_channels": 14,
					    "time_embedding_type": "positional",
					    "use_timestep_embedding": True,
					    "flip_sin_to_cos": False,
					    "freq_shift": 1.0,
					    "out_block_type": "OutConv1DBlock",
					    "mid_block_type": "MidResTemporalBlock1D",
					    "down_block_types": ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
					    "up_block_types": ("UpResnetBlock1D", "UpResnetBlock1D", "UpResnetBlock1D"),
					    "act_fn": "mish",
					}
					_UpperCAmelCase		:							List[Any]             =   self.dummy_input
					return init_dict, inputs_dict




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  List[Any]			) ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							int             =   UNetaDModel.from_pretrained(
					    "bglick13/hopper-medium-v2-value-function-hor32"						,	output_loading_info=lowerCamelCase__						,	subfolder="unet"			)
					self.assertIsNotNone(lowerCamelCase__			)
					self.assertEqual(len(loading_info["missing_keys"]			)						,	0			)

					model.to(lowerCamelCase__			)
					_UpperCAmelCase		:							Dict             =   model(**self.dummy_input			)

					assert image is not None, "Make sure output is not None"




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  List[str]			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:							str             =   UNetaDModel.from_pretrained("bglick13/hopper-medium-v2-value-function-hor32"						,	subfolder="unet"			)
					torch.manual_seed(0			)
					if torch.cuda.is_available():
							torch.cuda.manual_seed_all(0			)

					_UpperCAmelCase		:							str             =   model.config.in_channels
					_UpperCAmelCase		:							Optional[int]             =   16
					_UpperCAmelCase		:							List[str]             =   torch.randn((1, seq_len, num_features)			).permute(
					    0						,	2						,	1			)  # match original, we can update values and remove
					_UpperCAmelCase		:							Tuple             =   torch.full((num_features,)						,	0			)

					with torch.no_grad():
							_UpperCAmelCase		:							List[Any]             =   model(lowerCamelCase__						,	lowerCamelCase__			).sample.permute(0						,	2						,	1			)

					_UpperCAmelCase		:							Optional[int]             =   output[0, -3:, -3:].flatten()
					# fmt: off
					_UpperCAmelCase		:							Dict             =   torch.tensor([-2.1_3_7_1_7_2, 1.1_4_2_6_0_1_6, 0.3_6_8_8_6_8_7, -0.7_6_6_9_2_2, 0.7_3_0_3_1_4_6, 0.1_1_0_3_8_8_6_4, -0.4_7_6_0_6_3_3, 0.1_3_2_7_0_1_7_2, 0.0_2_5_9_1_3_4_8]			)
					# fmt: on
					self.assertTrue(torch.allclose(lowerCamelCase__						,	lowerCamelCase__						,	rtol=1E-3			)			)




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



					'''simple docstring'''







					pass






			@slow
			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   "harmonai/maestro-150k"
					_UpperCAmelCase		:							Tuple             =   UNetaDModel.from_pretrained(lowerCamelCase__						,	subfolder="unet"			)
					model.to(lowerCamelCase__			)

					_UpperCAmelCase		:							List[str]             =   6_55_36
					_UpperCAmelCase		:							Any             =   torch.sin(torch.arange(lowerCamelCase__			)[None, None, :].repeat(1						,	2						,	1			)			).to(lowerCamelCase__			)
					_UpperCAmelCase		:							List[Any]             =   torch.tensor([1]			).to(lowerCamelCase__			)

					with torch.no_grad():
							_UpperCAmelCase		:							List[str]             =   model(lowerCamelCase__						,	lowerCamelCase__			).sample

					_UpperCAmelCase		:							Union[str, Any]             =   output.abs().sum()
					_UpperCAmelCase		:							Any             =   output.abs().max()

					assert (output_sum - 2_2_4.0_8_9_6).abs() < 0.5
					assert (output_max - 0.0_6_0_7).abs() < 4E-4






class 						lowerCAmelCase__      (      UpperCAmelCase__     ,							UpperCAmelCase__     ,							unittest.TestCase						):
			lowerCAmelCase			:						Union[str, Any]							 =  UNetaDModel
			lowerCAmelCase			:						Union[str, Any]							 =  "sample"




			@property
			def 				lowerCAmelCase__     (		self	:  List[Any]			) ->str:



					'''simple docstring'''







					_UpperCAmelCase		:							Tuple             =   4
					_UpperCAmelCase		:							Optional[int]             =   14
					_UpperCAmelCase		:							Optional[int]             =   16

					_UpperCAmelCase		:							Optional[int]             =   floats_tensor((batch_size, num_features, seq_len)			).to(lowerCamelCase__			)
					_UpperCAmelCase		:							Union[str, Any]             =   torch.tensor([10] * batch_size			).to(lowerCamelCase__			)

					return {"sample": noise, "timestep": time_step}




			@property
			def 				lowerCAmelCase__     (		self	:  Dict			) ->Optional[int]:



					'''simple docstring'''







					return (4, 14, 16)




			@property
			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Optional[int]:



					'''simple docstring'''







					return (4, 14, 1)




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  List[str]			) ->Dict:



					'''simple docstring'''







					super().test_determinism()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  int			) ->Optional[Any]:



					'''simple docstring'''







					super().test_outputs_equivalence()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  str			) ->Optional[int]:



					'''simple docstring'''







					super().test_from_save_pretrained()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  int			) ->Dict:



					'''simple docstring'''







					super().test_from_save_pretrained_variant()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  int			) ->Dict:



					'''simple docstring'''







					super().test_model_from_pretrained()




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Dict:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							str             =   self.prepare_init_args_and_inputs_for_common()
					_UpperCAmelCase		:							Optional[int]             =   self.model_class(**lowerCamelCase__			)
					model.to(lowerCamelCase__			)
					model.eval()

					with torch.no_grad():
							_UpperCAmelCase		:							int             =   model(**lowerCamelCase__			)

							if isinstance(lowerCamelCase__						,	lowerCamelCase__			):
									_UpperCAmelCase		:							List[Any]             =   output.sample

					self.assertIsNotNone(lowerCamelCase__			)
					_UpperCAmelCase		:							Any             =   torch.Size((inputs_dict["sample"].shape[0], 1)			)
					self.assertEqual(output.shape						,	lowerCamelCase__						,	"Input and output shapes do not match"			)




			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Dict:



					'''simple docstring'''







					pass




			def 				lowerCAmelCase__     (		self	:  List[str]			) ->List[str]:



					'''simple docstring'''







					pass




			def 				lowerCAmelCase__     (		self	:  Dict			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:							List[str]             =   {
					    "in_channels": 14,
					    "out_channels": 14,
					    "down_block_types": ["DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"],
					    "up_block_types": [],
					    "out_block_type": "ValueFunction",
					    "mid_block_type": "ValueFunctionMidBlock1D",
					    "block_out_channels": [32, 64, 1_28, 2_56],
					    "layers_per_block": 1,
					    "downsample_each_block": True,
					    "use_timestep_embedding": True,
					    "freq_shift": 1.0,
					    "flip_sin_to_cos": False,
					    "time_embedding_type": "positional",
					    "act_fn": "mish",
					}
					_UpperCAmelCase		:							List[Any]             =   self.dummy_input
					return init_dict, inputs_dict




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  List[str]			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Any             =   UNetaDModel.from_pretrained(
					    "bglick13/hopper-medium-v2-value-function-hor32"						,	output_loading_info=lowerCamelCase__						,	subfolder="value_function"			)
					self.assertIsNotNone(lowerCamelCase__			)
					self.assertEqual(len(vf_loading_info["missing_keys"]			)						,	0			)

					value_function.to(lowerCamelCase__			)
					_UpperCAmelCase		:							Optional[int]             =   value_function(**self.dummy_input			)

					assert image is not None, "Make sure output is not None"




			@unittest.skipIf(torch_device == "mps"						,	"mish op not supported in MPS"			)
			def 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							List[str]             =   UNetaDModel.from_pretrained(
					    "bglick13/hopper-medium-v2-value-function-hor32"						,	output_loading_info=lowerCamelCase__						,	subfolder="value_function"			)
					torch.manual_seed(0			)
					if torch.cuda.is_available():
							torch.cuda.manual_seed_all(0			)

					_UpperCAmelCase		:							Tuple             =   value_function.config.in_channels
					_UpperCAmelCase		:							Tuple             =   14
					_UpperCAmelCase		:							int             =   torch.randn((1, seq_len, num_features)			).permute(
					    0						,	2						,	1			)  # match original, we can update values and remove
					_UpperCAmelCase		:							Tuple             =   torch.full((num_features,)						,	0			)

					with torch.no_grad():
							_UpperCAmelCase		:							int             =   value_function(lowerCamelCase__						,	lowerCamelCase__			).sample

					# fmt: off
					_UpperCAmelCase		:							Union[str, Any]             =   torch.tensor([1_6_5.2_5] * seq_len			)
					# fmt: on
					self.assertTrue(torch.allclose(lowerCamelCase__						,	lowerCamelCase__						,	rtol=1E-3			)			)






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



					'''simple docstring'''







					pass



