


'''simple docstring'''




import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()







class      _lowerCAmelCase							(      unittest.TestCase			):


	@property
	def    _a     (self							):
				torch.manual_seed(0							)
				A_     :						List[str]     =					UNetaDModel(
				    block_out_channels=(32, 64)	,				layers_per_block=2	,				sample_size=32	,				in_channels=3	,				out_channels=3	,				down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")	,				up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")	,				)
				return model


	@property
	def    _a     (self							):
				torch.manual_seed(0							)
				A_     :						Any     =					VQModel(
				    block_out_channels=[32, 64]	,				in_channels=3	,				out_channels=3	,				down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]	,				up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]	,				latent_channels=3	,				)
				return model


	@property
	def    _a     (self							):
				torch.manual_seed(0							)
				A_     :						Dict     =					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=1000	,				)
				return CLIPTextModel(_lowerCAmelCase							)


	def    _a     (self							):
				A_     :						List[Any]     =					self.dummy_uncond_unet
				A_     :						Dict     =					DDIMScheduler()
				A_     :						Any     =					self.dummy_vq_model

				A_     :						Any     =					LDMPipeline(unet=_lowerCAmelCase	,				vqvae=_lowerCAmelCase	,				scheduler=_lowerCAmelCase							)
				ldm.to(_lowerCAmelCase							)
				ldm.set_progress_bar_config(disable=_lowerCAmelCase							)

				A_     :						int     =					torch.manual_seed(0							)
				A_     :						Tuple     =					ldm(generator=_lowerCAmelCase	,				num_inference_steps=2	,				output_type="""numpy"""							).images

				A_     :						Optional[int]     =					torch.manual_seed(0							)
				A_     :						Optional[int]     =					ldm(generator=_lowerCAmelCase	,				num_inference_steps=2	,				output_type="""numpy"""	,				return_dict=_lowerCAmelCase							)[0]

				A_     :						Optional[Any]     =					image[0, -3:, -3:, -1]
				A_     :						Union[str, Any]     =					image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 64, 64, 3)
				A_     :						List[str]     =					np.array([0.85_12, 0.8_18, 0.64_11, 0.68_08, 0.44_65, 0.56_18, 0.46, 0.62_31, 0.51_72]							)
				A_     :						Any     =					1E-2 if torch_device != """mps""" else 3E-2

				assert np.abs(image_slice.flatten() - expected_slice							).max() < tolerance
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < tolerance







@slow
@require_torch
class      _lowerCAmelCase							(      unittest.TestCase			):


	def    _a     (self							):
				A_     :						List[Any]     =					LDMPipeline.from_pretrained("""CompVis/ldm-celebahq-256"""							)
				ldm.to(_lowerCAmelCase							)
				ldm.set_progress_bar_config(disable=_lowerCAmelCase							)

				A_     :						List[Any]     =					torch.manual_seed(0							)
				A_     :						Any     =					ldm(generator=_lowerCAmelCase	,				num_inference_steps=5	,				output_type="""numpy"""							).images

				A_     :						List[Any]     =					image[0, -3:, -3:, -1]

				assert image.shape == (1, 256, 256, 3)
				A_     :						Optional[Any]     =					np.array([0.43_99, 0.4_49_75, 0.4_68_25, 0.4_74, 0.43_59, 0.45_81, 0.4_50_95, 0.43_41, 0.44_47]							)
				A_     :						Tuple     =					1E-2 if torch_device != """mps""" else 3E-2

				assert np.abs(image_slice.flatten() - expected_slice							).max() < tolerance




'''simple docstring'''




import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch







@require_torch
class      _lowerCAmelCase							:


	@property
	def    _a     (self							):
				return self.get_dummy_input()


	@property
	def    _a     (self							):
				if self.block_type == "down":
							return (4, 32, 16, 16)
				elif self.block_type == "mid":
							return (4, 32, 32, 32)
				elif self.block_type == "up":
							return (4, 32, 64, 64)

				raise ValueError(F'\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'							)


	def    _a     (self	,				lowercase=True	,				lowercase=False	,				lowercase=False	,				lowercase=False	,				):
				A_     :						List[str]     =					4
				A_     :						int     =					32
				A_     :						Optional[int]     =					(32, 32)

				A_     :						Optional[Any]     =					torch.manual_seed(0							)
				A_     :						int     =					torch.device(lowercase							)
				A_     :						int     =					(batch_size, num_channels) + sizes
				A_     :						Optional[int]     =					randn_tensor(lowercase	,				generator=lowercase	,				device=lowercase							)
				A_     :						Union[str, Any]     =					{"""hidden_states""": hidden_states}

				if include_temb:
							A_     :						str     =					128
							A_     :						List[Any]     =					randn_tensor((batch_size, temb_channels)	,				generator=lowercase	,				device=lowercase							)

				if include_res_hidden_states_tuple:
							A_     :						List[str]     =					torch.manual_seed(1							)
							A_     :						int     =					(randn_tensor(lowercase	,				generator=lowercase	,				device=lowercase							),)

				if include_encoder_hidden_states:
							A_     :						List[str]     =					floats_tensor((batch_size, 32, 32)							).to(lowercase							)

				if include_skip_sample:
							A_     :						Dict     =					randn_tensor(((batch_size, 3) + sizes)	,				generator=lowercase	,				device=lowercase							)

				return dummy_input


	def    _a     (self							):
				A_     :						Tuple     =					{
				    """in_channels""": 32,
				    """out_channels""": 32,
				    """temb_channels""": 128,
				}
				if self.block_type == "up":
							A_     :						Any     =					32

				if self.block_type == "mid":
							init_dict.pop("""out_channels"""							)

				A_     :						Optional[int]     =					self.dummy_input
				return init_dict, inputs_dict


	def    _a     (self	,				lowercase							):
				A_,   A_     :						Optional[Any]     =					self.prepare_init_args_and_inputs_for_common()
				A_     :						int     =					self.block_class(**lowercase							)
				unet_block.to(lowercase							)
				unet_block.eval()

				with torch.no_grad():
							A_     :						List[str]     =					unet_block(**lowercase							)

				if isinstance(lowercase	,				lowercase							):
							A_     :						Union[str, Any]     =					output[0]

				self.assertEqual(output.shape	,				self.output_shape							)

				A_     :						int     =					output[0, -1, -3:, -3:]
				A_     :						List[Any]     =					torch.tensor(lowercase							).to(lowercase							)
				assert torch_all_close(output_slice.flatten()	,				lowercase	,				atol=5E-3							)


	@unittest.skipIf(torch_device == """mps"""	,				"""Training is not supported in mps"""							)
	def    _a     (self							):
				A_,   A_     :						Tuple     =					self.prepare_init_args_and_inputs_for_common()
				A_     :						List[str]     =					self.block_class(**lowercase							)
				model.to(lowercase							)
				model.train()
				A_     :						Any     =					model(**lowercase							)

				if isinstance(lowercase	,				lowercase							):
							A_     :						str     =					output[0]

				A_     :						Union[str, Any]     =					torch.device(lowercase							)
				A_     :						Tuple     =					randn_tensor(output.shape	,				device=lowercase							)
				A_     :						List[str]     =					torch.nn.functional.mse_loss(lowercase	,				lowercase							)
				loss.backward()