


"""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 			UpperCamelCase_     (			lowerCamelCase__      ,  lowerCamelCase__      ,  lowerCamelCase__      ,  unittest.TestCase):

	"""simple docstring"""





	snake_case__    :  str    =      AltDiffusionPipeline
	snake_case__    :  Dict    =      TEXT_TO_IMAGE_PARAMS
	snake_case__    :  Dict    =      TEXT_TO_IMAGE_BATCH_PARAMS
	snake_case__    :  Optional[Any]    =      TEXT_TO_IMAGE_IMAGE_PARAMS
	snake_case__    :  Union[str, Any]    =      TEXT_TO_IMAGE_IMAGE_PARAMS




	def   UpperCAmelCase_				(			self      :    int       )	->     int:
						torch.manual_seed(0       )
						__SCREAMING_SNAKE_CASE						       =    UNetaDConditionModel(
						    block_out_channels=(3_2, 6_4)      ,      layers_per_block=2      ,      sample_size=3_2      ,      in_channels=4      ,      out_channels=4      ,      down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")      ,      up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")      ,      cross_attention_dim=3_2      ,      )
						__SCREAMING_SNAKE_CASE						       =    DDIMScheduler(
						    beta_start=0.00_085      ,      beta_end=0.012      ,      beta_schedule="scaled_linear"      ,      clip_sample=lowercase__      ,      set_alpha_to_one=lowercase__      ,      )
						torch.manual_seed(0       )
						__SCREAMING_SNAKE_CASE						       =    AutoencoderKL(
						    block_out_channels=[3_2, 6_4]      ,      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       )
						__SCREAMING_SNAKE_CASE						       =    CLIPTextConfig(
						    bos_token_id=0      ,      eos_token_id=2      ,      hidden_size=3_2      ,      projection_dim=3_2      ,      intermediate_size=3_7      ,      layer_norm_eps=1E-05      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      pad_token_id=1      ,      vocab_size=5_0_0_2      ,      )
						__SCREAMING_SNAKE_CASE						       =    CLIPTextModel(lowercase__       )

						__SCREAMING_SNAKE_CASE						       =    XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"       )
						__SCREAMING_SNAKE_CASE						       =    7_7

						__SCREAMING_SNAKE_CASE						       =    {
						    "unet": unet,
						    "scheduler": scheduler,
						    "vae": vae,
						    "text_encoder": text_encoder,
						    "tokenizer": tokenizer,
						    "safety_checker": None,
						    "feature_extractor": None,
						}
						return components




	def   UpperCAmelCase_				(			self      :    int      ,      UpperCAmelCase__      :    Optional[Any]      ,      UpperCAmelCase__      :    Optional[int]=0       )	->     int:
						if str(lowercase__       ).startswith("mps"       ):
											__SCREAMING_SNAKE_CASE						       =    torch.manual_seed(lowercase__       )
						else:
											__SCREAMING_SNAKE_CASE						       =    torch.Generator(device=lowercase__       ).manual_seed(lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    {
						    "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   UpperCAmelCase_				(			self      :    Optional[int]       )	->     Optional[int]:
						super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )




	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Optional[Any]:
						super().test_inference_batch_single_identical(expected_max_diff=3E-3       )




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     Union[str, Any]:
						__SCREAMING_SNAKE_CASE						       =    "cpu"  # ensure determinism for the device-dependent torch.Generator

						__SCREAMING_SNAKE_CASE						       =    self.get_dummy_components()
						torch.manual_seed(0       )
						__SCREAMING_SNAKE_CASE						       =    RobertaSeriesConfig(
						    hidden_size=3_2      ,      project_dim=3_2      ,      intermediate_size=3_7      ,      layer_norm_eps=1E-05      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      vocab_size=5_0_0_2      ,      )
						# TODO: remove after fixing the non-deterministic text encoder
						__SCREAMING_SNAKE_CASE						       =    RobertaSeriesModelWithTransformation(lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    text_encoder

						__SCREAMING_SNAKE_CASE						       =    AltDiffusionPipeline(**lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    alt_pipe.to(lowercase__       )
						alt_pipe.set_progress_bar_config(disable=lowercase__       )

						__SCREAMING_SNAKE_CASE						       =    self.get_dummy_inputs(lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    "A photo of an astronaut"
						__SCREAMING_SNAKE_CASE						       =    alt_pipe(**lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    output.images
						__SCREAMING_SNAKE_CASE						       =    image[0, -3:, -3:, -1]

						assert image.shape == (1, 6_4, 6_4, 3)
						__SCREAMING_SNAKE_CASE						       =    np.array(
						    [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093]       )

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




	def   UpperCAmelCase_				(			self      :    Optional[int]       )	->     Any:
						__SCREAMING_SNAKE_CASE						       =    "cpu"  # ensure determinism for the device-dependent torch.Generator

						__SCREAMING_SNAKE_CASE						       =    self.get_dummy_components()
						__SCREAMING_SNAKE_CASE						       =    PNDMScheduler(skip_prk_steps=lowercase__       )
						torch.manual_seed(0       )
						__SCREAMING_SNAKE_CASE						       =    RobertaSeriesConfig(
						    hidden_size=3_2      ,      project_dim=3_2      ,      intermediate_size=3_7      ,      layer_norm_eps=1E-05      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      vocab_size=5_0_0_2      ,      )
						# TODO: remove after fixing the non-deterministic text encoder
						__SCREAMING_SNAKE_CASE						       =    RobertaSeriesModelWithTransformation(lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    text_encoder
						__SCREAMING_SNAKE_CASE						       =    AltDiffusionPipeline(**lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    alt_pipe.to(lowercase__       )
						alt_pipe.set_progress_bar_config(disable=lowercase__       )

						__SCREAMING_SNAKE_CASE						       =    self.get_dummy_inputs(lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    alt_pipe(**lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    output.images
						__SCREAMING_SNAKE_CASE						       =    image[0, -3:, -3:, -1]

						assert image.shape == (1, 6_4, 6_4, 3)
						__SCREAMING_SNAKE_CASE						       =    np.array(
						    [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237]       )

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




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

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Tuple:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def   UpperCAmelCase_				(			self      :    Any       )	->     str:
						# make sure here that pndm scheduler skips prk
						__SCREAMING_SNAKE_CASE						       =    AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"      ,      safety_checker=lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    alt_pipe.to(lowercase__       )
						alt_pipe.set_progress_bar_config(disable=lowercase__       )

						__SCREAMING_SNAKE_CASE						       =    "A painting of a squirrel eating a burger"
						__SCREAMING_SNAKE_CASE						       =    torch.manual_seed(0       )
						__SCREAMING_SNAKE_CASE						       =    alt_pipe([prompt]      ,      generator=lowercase__      ,      guidance_scale=6.0      ,      num_inference_steps=2_0      ,      output_type="np"       )

						__SCREAMING_SNAKE_CASE						       =    output.images

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						__SCREAMING_SNAKE_CASE						       =    np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586]       )

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




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Optional[int]:
						__SCREAMING_SNAKE_CASE						       =    DDIMScheduler.from_pretrained("BAAI/AltDiffusion"      ,      subfolder="scheduler"       )

						__SCREAMING_SNAKE_CASE						       =    AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"      ,      scheduler=lowercase__      ,      safety_checker=lowercase__       )
						__SCREAMING_SNAKE_CASE						       =    alt_pipe.to(lowercase__       )
						alt_pipe.set_progress_bar_config(disable=lowercase__       )

						__SCREAMING_SNAKE_CASE						       =    "A painting of a squirrel eating a burger"
						__SCREAMING_SNAKE_CASE						       =    torch.manual_seed(0       )

						__SCREAMING_SNAKE_CASE						       =    alt_pipe([prompt]      ,      generator=lowercase__      ,      num_inference_steps=2      ,      output_type="numpy"       )
						__SCREAMING_SNAKE_CASE						       =    output.images

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						__SCREAMING_SNAKE_CASE						       =    np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323]       )

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




"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import _LazyModule


a__    :							Optional[int]            =					{'''tokenization_byt5''': ['''ByT5Tokenizer''']}


if TYPE_CHECKING:
		from .tokenization_byta import ByTaTokenizer
else:
		import sys

		a__    :							Any            =					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
