
'''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  lowerCAmelCase__   (  lowerCamelCase_    ,						lowerCamelCase_    ,						lowerCamelCase_    ,						unittest.TestCase     ):
					lowerCAmelCase_        =     AltDiffusionPipeline
					lowerCAmelCase_        =     TEXT_TO_IMAGE_PARAMS
					lowerCAmelCase_        =     TEXT_TO_IMAGE_BATCH_PARAMS
					lowerCAmelCase_        =     TEXT_TO_IMAGE_IMAGE_PARAMS
					lowerCAmelCase_        =     TEXT_TO_IMAGE_IMAGE_PARAMS






					def 	_snake_case    (   self		):


											"""simple docstring"""





											torch.manual_seed(0		)
											lowercase_	:     Optional[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				,							)
											lowercase_	:     Dict													=       DDIMScheduler(
											    beta_start=0.00_085				,							beta_end=0.012				,							beta_schedule='''scaled_linear'''				,							clip_sample=__SCREAMING_SNAKE_CASE				,							set_alpha_to_one=__SCREAMING_SNAKE_CASE				,							)
											torch.manual_seed(0		)
											lowercase_	:     str													=       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		)
											lowercase_	:     Union[str, Any]													=       CLIPTextConfig(
											    bos_token_id=0				,							eos_token_id=2				,							hidden_size=32				,							projection_dim=32				,							intermediate_size=37				,							layer_norm_eps=1E-0_5				,							num_attention_heads=4				,							num_hidden_layers=5				,							pad_token_id=1				,							vocab_size=50_02				,							)
											lowercase_	:     Dict													=       CLIPTextModel(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Union[str, Any]													=       XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
											lowercase_	:     Union[str, Any]													=       77

											lowercase_	:     Dict													=       {
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''vae''': vae,
											    '''text_encoder''': text_encoder,
											    '''tokenizer''': tokenizer,
											    '''safety_checker''': None,
											    '''feature_extractor''': None,
											}
											return components






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE=0		):


											"""simple docstring"""





											if str(__SCREAMING_SNAKE_CASE		).startswith('''mps'''		):
																	lowercase_	:     Optional[int]													=       torch.manual_seed(__SCREAMING_SNAKE_CASE		)
											else:
																	lowercase_	:     Optional[Any]													=       torch.Generator(device=__SCREAMING_SNAKE_CASE		).manual_seed(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     List[str]													=       {
											    '''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 	_snake_case    (   self		):


											"""simple docstring"""





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






					def 	_snake_case    (   self		):


											"""simple docstring"""





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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Union[str, Any]													=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator

											lowercase_	:     List[Any]													=       self.get_dummy_components()
											torch.manual_seed(0		)
											lowercase_	:     Optional[int]													=       RobertaSeriesConfig(
											    hidden_size=32				,							project_dim=32				,							intermediate_size=37				,							layer_norm_eps=1E-0_5				,							num_attention_heads=4				,							num_hidden_layers=5				,							vocab_size=50_02				,							)
											# TODO: remove after fixing the non-deterministic text encoder
											lowercase_	:     Union[str, Any]													=       RobertaSeriesModelWithTransformation(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     List[Any]													=       text_encoder

											lowercase_	:     int													=       AltDiffusionPipeline(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Union[str, Any]													=       alt_pipe.to(__SCREAMING_SNAKE_CASE		)
											alt_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[Any]													=       self.get_dummy_inputs(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[int]													=       '''A photo of an astronaut'''
											lowercase_	:     Optional[int]													=       alt_pipe(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[Any]													=       output.images
											lowercase_	:     Union[str, Any]													=       image[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											lowercase_	:     List[str]													=       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 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     int													=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator

											lowercase_	:     int													=       self.get_dummy_components()
											lowercase_	:     Optional[int]													=       PNDMScheduler(skip_prk_steps=__SCREAMING_SNAKE_CASE		)
											torch.manual_seed(0		)
											lowercase_	:     str													=       RobertaSeriesConfig(
											    hidden_size=32				,							project_dim=32				,							intermediate_size=37				,							layer_norm_eps=1E-0_5				,							num_attention_heads=4				,							num_hidden_layers=5				,							vocab_size=50_02				,							)
											# TODO: remove after fixing the non-deterministic text encoder
											lowercase_	:     str													=       RobertaSeriesModelWithTransformation(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[Any]													=       text_encoder
											lowercase_	:     List[str]													=       AltDiffusionPipeline(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[Any]													=       alt_pipe.to(__SCREAMING_SNAKE_CASE		)
											alt_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE		)

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

											assert image.shape == (1, 64, 64, 3)
											lowercase_	:     Optional[Any]													=       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  lowerCAmelCase__   (  unittest.TestCase     ):






					def 	_snake_case    (   self		):


											"""simple docstring"""





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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Dict													=       AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''				,							safety_checker=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Dict													=       alt_pipe.to(__SCREAMING_SNAKE_CASE		)
											alt_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[Any]													=       '''A painting of a squirrel eating a burger'''
											lowercase_	:     int													=       torch.manual_seed(0		)
											lowercase_	:     Optional[Any]													=       alt_pipe([prompt]				,							generator=__SCREAMING_SNAKE_CASE				,							guidance_scale=6.0				,							num_inference_steps=20				,							output_type='''np'''		)

											lowercase_	:     Tuple													=       output.images

											lowercase_	:     int													=       image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											lowercase_	:     List[str]													=       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 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''				,							subfolder='''scheduler'''		)

											lowercase_	:     Optional[Any]													=       AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''				,							scheduler=__SCREAMING_SNAKE_CASE				,							safety_checker=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Union[str, Any]													=       alt_pipe.to(__SCREAMING_SNAKE_CASE		)
											alt_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[Any]													=       '''A painting of a squirrel eating a burger'''
											lowercase_	:     Optional[int]													=       torch.manual_seed(0		)

											lowercase_	:     Optional[Any]													=       alt_pipe([prompt]				,							generator=__SCREAMING_SNAKE_CASE				,							num_inference_steps=2				,							output_type='''numpy'''		)
											lowercase_	:     Tuple													=       output.images

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

											assert image.shape == (1, 5_12, 5_12, 3)
											lowercase_	:     str													=       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 OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_lowercase :  Union[str, Any]    							=       {
    "configuration_mask2former": [
        "MASK2FORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "Mask2FormerConfig",
    ],
}

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							_lowercase :  Optional[int]    							=       ["Mask2FormerImageProcessor"]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							_lowercase :  str    							=       [
							    "MASK2FORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "Mask2FormerForUniversalSegmentation",
							    "Mask2FormerModel",
							    "Mask2FormerPreTrainedModel",
							]

if TYPE_CHECKING:
							from .configuration_maskaformer import MASK2FORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, MaskaFormerConfig

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .image_processing_maskaformer import MaskaFormerImageProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_maskaformer import (
														    MASK2FORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
														    MaskaFormerForUniversalSegmentation,
														    MaskaFormerModel,
														    MaskaFormerPreTrainedModel,
														)


else:
							import sys

							_lowercase :  Dict    							=       _LazyModule(__name__, globals()["__file__"], _import_structure)



