



"""simple docstring"""






import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class     A__     (						__SCREAMING_SNAKE_CASE						,	__SCREAMING_SNAKE_CASE						,	__SCREAMING_SNAKE_CASE						,	unittest.TestCase							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE					   =			AltDiffusionPipeline
	SCREAMING_SNAKE_CASE					   =			TEXT_TO_IMAGE_PARAMS
	SCREAMING_SNAKE_CASE					   =			TEXT_TO_IMAGE_BATCH_PARAMS
	SCREAMING_SNAKE_CASE					   =			TEXT_TO_IMAGE_IMAGE_PARAMS
	SCREAMING_SNAKE_CASE					   =			TEXT_TO_IMAGE_IMAGE_PARAMS

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[Any])       -> Tuple:

				"""simple docstring"""

				torch.manual_seed(0)
				__lowerCAmelCase   :						Optional[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 ,    )
				__lowerCAmelCase   :						Union[str, Any]				= DDIMScheduler(
				    beta_start=0.0_0085 ,    beta_end=0.012 ,    beta_schedule="scaled_linear" ,    clip_sample=_SCREAMING_SNAKE_CASE ,    set_alpha_to_one=_SCREAMING_SNAKE_CASE ,    )
				torch.manual_seed(0)
				__lowerCAmelCase   :						Dict				= 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 ,    )

				# TODO: address the non-deterministic text encoder (fails for save-load tests)
				# torch.manual_seed(0)
				# text_encoder_config = RobertaSeriesConfig(
				#     hidden_size=32,
				#     project_dim=32,
				#     intermediate_size=37,
				#     layer_norm_eps=1e-05,
				#     num_attention_heads=4,
				#     num_hidden_layers=5,
				#     vocab_size=5002,
				# )
				# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

				torch.manual_seed(0)
				__lowerCAmelCase   :						str				= CLIPTextConfig(
				    bos_token_id=0 ,    eos_token_id=2 ,    hidden_size=32 ,    projection_dim=32 ,    intermediate_size=37 ,    layer_norm_eps=1e-05 ,    num_attention_heads=4 ,    num_hidden_layers=5 ,    pad_token_id=1 ,    vocab_size=5002 ,    )
				__lowerCAmelCase   :						Tuple				= CLIPTextModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Any				= XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta")
				__lowerCAmelCase   :						Optional[Any]				= 77

				__lowerCAmelCase   :						Union[str, Any]				= {
				    "unet": unet,
				    "scheduler": scheduler,
				    "vae": vae,
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
				return components

	def 							_SCREAMING_SNAKE_CASE     (  self:	str ,    _SCREAMING_SNAKE_CASE:	Union[str, Any] ,    _SCREAMING_SNAKE_CASE:	str=0)       -> List[Any]:

				"""simple docstring"""

				if str(_SCREAMING_SNAKE_CASE).startswith("mps"):
							__lowerCAmelCase   :						Any				= torch.manual_seed(_SCREAMING_SNAKE_CASE)
				else:
							__lowerCAmelCase   :						Union[str, Any]				= torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Any				= {
				    "prompt": "A painting of a squirrel eating a burger",
				    "generator": generator,
				    "num_inference_steps": 2,
				    "guidance_scale": 6.0,
				    "output_type": "numpy",
				}
				return inputs

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

				"""simple docstring"""

				super().test_attention_slicing_forward_pass(expected_max_diff=3e-3)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Union[str, Any])       -> Tuple:

				"""simple docstring"""

				super().test_inference_batch_single_identical(expected_max_diff=3e-3)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[int])       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Dict				= "cpu"  # ensure determinism for the device-dependent torch.Generator

				__lowerCAmelCase   :						int				= self.get_dummy_components()
				torch.manual_seed(0)
				__lowerCAmelCase   :						List[str]				= RobertaSeriesConfig(
				    hidden_size=32 ,    project_dim=32 ,    intermediate_size=37 ,    layer_norm_eps=1e-05 ,    num_attention_heads=4 ,    num_hidden_layers=5 ,    vocab_size=5002 ,    )
				# TODO: remove after fixing the non-deterministic text encoder
				__lowerCAmelCase   :						List[Any]				= RobertaSeriesModelWithTransformation(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						int				= text_encoder

				__lowerCAmelCase   :						List[Any]				= AltDiffusionPipeline(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Optional[Any]				= alt_pipe.to(_SCREAMING_SNAKE_CASE)
				alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Tuple				= self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Any				= "A photo of an astronaut"
				__lowerCAmelCase   :						Optional[Any]				= alt_pipe(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Optional[Any]				= output.images
				__lowerCAmelCase   :						List[Any]				= image[0, -3:, -3:, -1]

				assert image.shape == (1, 64, 64, 3)
				__lowerCAmelCase   :						List[Any]				= np.array(
				    [0.574_8162, 0.6044_7145, 0.4882_1217, 0.5010_0636, 0.543_1185, 0.4576_3683, 0.4965_7696, 0.4813_2733, 0.4757_3093])

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







	def 							_SCREAMING_SNAKE_CASE     (  self:	Tuple)       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						int				= "cpu"  # ensure determinism for the device-dependent torch.Generator

				__lowerCAmelCase   :						Tuple				= self.get_dummy_components()
				__lowerCAmelCase   :						List[Any]				= PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE)
				torch.manual_seed(0)
				__lowerCAmelCase   :						Union[str, Any]				= RobertaSeriesConfig(
				    hidden_size=32 ,    project_dim=32 ,    intermediate_size=37 ,    layer_norm_eps=1e-05 ,    num_attention_heads=4 ,    num_hidden_layers=5 ,    vocab_size=5002 ,    )
				# TODO: remove after fixing the non-deterministic text encoder
				__lowerCAmelCase   :						Any				= RobertaSeriesModelWithTransformation(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Tuple				= text_encoder
				__lowerCAmelCase   :						Dict				= AltDiffusionPipeline(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= alt_pipe.to(_SCREAMING_SNAKE_CASE)
				alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						int				= self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Optional[int]				= alt_pipe(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						List[Any]				= output.images
				__lowerCAmelCase   :						List[str]				= image[0, -3:, -3:, -1]

				assert image.shape == (1, 64, 64, 3)
				__lowerCAmelCase   :						Dict				= np.array(
				    [0.5160_5093, 0.570_7241, 0.4736_5507, 0.5057_8886, 0.563_3877, 0.464_2503, 0.518_2081, 0.4876_3484, 0.4908_4237])

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






@slow
@require_torch_gpu
class     A__     (						unittest.TestCase							):




	'''simple docstring'''





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

				"""simple docstring"""

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

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> Union[str, Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						str				= AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion" ,    safety_checker=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Dict				= alt_pipe.to(_SCREAMING_SNAKE_CASE)
				alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						int				= "A painting of a squirrel eating a burger"
				__lowerCAmelCase   :						List[Any]				= torch.manual_seed(0)
				__lowerCAmelCase   :						Union[str, Any]				= alt_pipe([prompt] ,    generator=_SCREAMING_SNAKE_CASE ,    guidance_scale=6.0 ,    num_inference_steps=20 ,    output_type="np")

				__lowerCAmelCase   :						str				= output.images

				__lowerCAmelCase   :						List[str]				= image[0, -3:, -3:, -1]

				assert image.shape == (1, 512, 512, 3)
				__lowerCAmelCase   :						Union[str, Any]				= np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586])

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







	def 							_SCREAMING_SNAKE_CASE     (  self:	int)       -> Dict:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[Any]				= DDIMScheduler.from_pretrained("BAAI/AltDiffusion" ,    subfolder="scheduler")

				__lowerCAmelCase   :						Union[str, Any]				= AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion" ,    scheduler=_SCREAMING_SNAKE_CASE ,    safety_checker=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						List[str]				= alt_pipe.to(_SCREAMING_SNAKE_CASE)
				alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Dict				= "A painting of a squirrel eating a burger"
				__lowerCAmelCase   :						List[str]				= torch.manual_seed(0)

				__lowerCAmelCase   :						Tuple				= alt_pipe([prompt] ,    generator=_SCREAMING_SNAKE_CASE ,    num_inference_steps=2 ,    output_type="numpy")
				__lowerCAmelCase   :						Union[str, Any]				= output.images

				__lowerCAmelCase   :						Union[str, Any]				= image[0, -3:, -3:, -1]

				assert image.shape == (1, 512, 512, 3)
				__lowerCAmelCase   :						Optional[int]				= np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323])

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




"""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     A__     :




	'''simple docstring'''





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

				"""simple docstring"""

				return self.get_dummy_input()

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

				"""simple docstring"""

				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 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    _SCREAMING_SNAKE_CASE:	Union[str, Any]=True ,    _SCREAMING_SNAKE_CASE:	Tuple=False ,    _SCREAMING_SNAKE_CASE:	Union[str, Any]=False ,    _SCREAMING_SNAKE_CASE:	Dict=False ,    )       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= 4
				__lowerCAmelCase   :						List[str]				= 32
				__lowerCAmelCase   :						str				= (32, 32)

				__lowerCAmelCase   :						int				= torch.manual_seed(0)
				__lowerCAmelCase   :						Dict				= torch.device(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= (batch_size, num_channels) + sizes
				__lowerCAmelCase   :						Dict				= randn_tensor(_SCREAMING_SNAKE_CASE ,    generator=_SCREAMING_SNAKE_CASE ,    device=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= {"hidden_states": hidden_states}

				if include_temb:
							__lowerCAmelCase   :						List[str]				= 128
							__lowerCAmelCase   :						Dict				= randn_tensor((batch_size, temb_channels) ,    generator=_SCREAMING_SNAKE_CASE ,    device=_SCREAMING_SNAKE_CASE)

				if include_res_hidden_states_tuple:
							__lowerCAmelCase   :						Dict				= torch.manual_seed(1)
							__lowerCAmelCase   :						int				= (randn_tensor(_SCREAMING_SNAKE_CASE ,    generator=_SCREAMING_SNAKE_CASE ,    device=_SCREAMING_SNAKE_CASE),)

				if include_encoder_hidden_states:
							__lowerCAmelCase   :						int				= floats_tensor((batch_size, 32, 32)).to(_SCREAMING_SNAKE_CASE)

				if include_skip_sample:
							__lowerCAmelCase   :						Optional[Any]				= randn_tensor(((batch_size, 3) + sizes) ,    generator=_SCREAMING_SNAKE_CASE ,    device=_SCREAMING_SNAKE_CASE)

				return dummy_input

	def 							_SCREAMING_SNAKE_CASE     (  self:	Tuple)       -> Any:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[int]				= {
				    "in_channels": 32,
				    "out_channels": 32,
				    "temb_channels": 128,
				}
				if self.block_type == "up":
							__lowerCAmelCase   :						Tuple				= 32

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

				__lowerCAmelCase   :						str				= self.dummy_input
				return init_dict, inputs_dict

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[str] ,    _SCREAMING_SNAKE_CASE:	List[str])       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= self.prepare_init_args_and_inputs_for_common()
				__lowerCAmelCase   :						Dict				= self.block_class(**_SCREAMING_SNAKE_CASE)
				unet_block.to(_SCREAMING_SNAKE_CASE)
				unet_block.eval()

				with torch.no_grad():
							__lowerCAmelCase   :						int				= unet_block(**_SCREAMING_SNAKE_CASE)

				if isinstance(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE):
							__lowerCAmelCase   :						List[Any]				= output[0]

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

				__lowerCAmelCase   :						str				= output[0, -1, -3:, -3:]
				__lowerCAmelCase   :						List[Any]				= torch.tensor(_SCREAMING_SNAKE_CASE).to(_SCREAMING_SNAKE_CASE)
				assert torch_all_close(output_slice.flatten() ,    _SCREAMING_SNAKE_CASE ,    atol=5e-3)







	@unittest.skipIf(torch_device == "mps" ,    "Training is not supported in mps")
	def 							_SCREAMING_SNAKE_CASE     (  self:	int)       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase		,			__lowerCAmelCase   :						Optional[Any]				= self.prepare_init_args_and_inputs_for_common()
				__lowerCAmelCase   :						str				= self.block_class(**_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.train()
				__lowerCAmelCase   :						Union[str, Any]				= model(**_SCREAMING_SNAKE_CASE)

				if isinstance(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE):
							__lowerCAmelCase   :						Union[str, Any]				= output[0]

				__lowerCAmelCase   :						List[Any]				= torch.device(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						int				= randn_tensor(output.shape ,    device=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						int				= torch.nn.functional.mse_loss(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
				loss.backward()