






"""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__			(    __magic_name__							,						__magic_name__							,						__magic_name__							,						unittest.TestCase   ):
   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       __a							(		self      :	Tuple							):


        '''simple docstring'''
        torch.manual_seed(0							)
        UpperCAmelCase__						:    Tuple							=							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							,			)
        UpperCAmelCase__						:    int							=							DDIMScheduler(
            beta_start=0.0_0085							,			beta_end=0.012							,			beta_schedule="scaled_linear"							,			clip_sample=snake_case__							,			set_alpha_to_one=snake_case__							,			)
        torch.manual_seed(0							)
        UpperCAmelCase__						:    Union[str, Any]							=							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							)
        UpperCAmelCase__						:    int							=							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							,			)
        UpperCAmelCase__						:    List[Any]							=							CLIPTextModel(snake_case__							)

        UpperCAmelCase__						:    List[Any]							=							XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"							)
        UpperCAmelCase__						:    str							=							7_7

        UpperCAmelCase__						:    Any							=							{
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components
   def       __a							(		self      :	List[str]							,			snake_case__      :	Optional[int]							,			snake_case__      :	Union[str, Any]=0							):


        '''simple docstring'''
        if str(snake_case__							).startswith("mps"							):
             UpperCAmelCase__						:    Any							=							torch.manual_seed(snake_case__							)
        else:
             UpperCAmelCase__						:    Any							=							torch.Generator(device=snake_case__							).manual_seed(snake_case__							)
        UpperCAmelCase__						:    int							=							{
            "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       __a							(		self      :	Tuple							):


        '''simple docstring'''
        super().test_attention_slicing_forward_pass(expected_max_diff=3e-3							)
   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        super().test_inference_batch_single_identical(expected_max_diff=3e-3							)
   def       __a							(		self      :	str							):


        '''simple docstring'''
        UpperCAmelCase__						:    Dict							=							"cpu"  # ensure determinism for the device-dependent torch.Generator

        UpperCAmelCase__						:    Dict							=							self.get_dummy_components()
        torch.manual_seed(0							)
        UpperCAmelCase__						:    Optional[Any]							=							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
        UpperCAmelCase__						:    Dict							=							RobertaSeriesModelWithTransformation(snake_case__							)
        UpperCAmelCase__						:    Any							=							text_encoder

        UpperCAmelCase__						:    Optional[int]							=							AltDiffusionPipeline(**snake_case__							)
        UpperCAmelCase__						:    List[str]							=							alt_pipe.to(snake_case__							)
        alt_pipe.set_progress_bar_config(disable=snake_case__							)

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

        assert image.shape == (1, 6_4, 6_4, 3)
        UpperCAmelCase__						:    Tuple							=							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       __a							(		self      :	Dict							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[int]							=							"cpu"  # ensure determinism for the device-dependent torch.Generator

        UpperCAmelCase__						:    Tuple							=							self.get_dummy_components()
        UpperCAmelCase__						:    Dict							=							PNDMScheduler(skip_prk_steps=snake_case__							)
        torch.manual_seed(0							)
        UpperCAmelCase__						:    Optional[int]							=							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
        UpperCAmelCase__						:    Optional[int]							=							RobertaSeriesModelWithTransformation(snake_case__							)
        UpperCAmelCase__						:    str							=							text_encoder
        UpperCAmelCase__						:    Tuple							=							AltDiffusionPipeline(**snake_case__							)
        UpperCAmelCase__						:    Tuple							=							alt_pipe.to(snake_case__							)
        alt_pipe.set_progress_bar_config(disable=snake_case__							)

        UpperCAmelCase__						:    Dict							=							self.get_dummy_inputs(snake_case__							)
        UpperCAmelCase__						:    Tuple							=							alt_pipe(**snake_case__							)
        UpperCAmelCase__						:    Optional[Any]							=							output.images
        UpperCAmelCase__						:    List[str]							=							image[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)
        UpperCAmelCase__						:    Any							=							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 			lowerCAmelCase__			(    unittest.TestCase   ):
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        # make sure here that pndm scheduler skips prk
        UpperCAmelCase__						:    Optional[int]							=							AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"							,			safety_checker=snake_case__							)
        UpperCAmelCase__						:    List[Any]							=							alt_pipe.to(snake_case__							)
        alt_pipe.set_progress_bar_config(disable=snake_case__							)

        UpperCAmelCase__						:    Optional[Any]							=							"A painting of a squirrel eating a burger"
        UpperCAmelCase__						:    Union[str, Any]							=							torch.manual_seed(0							)
        UpperCAmelCase__						:    Union[str, Any]							=							alt_pipe([prompt]							,			generator=snake_case__							,			guidance_scale=6.0							,			num_inference_steps=2_0							,			output_type="np"							)

        UpperCAmelCase__						:    List[str]							=							output.images

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

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        UpperCAmelCase__						:    Optional[int]							=							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       __a							(		self      :	Dict							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							DDIMScheduler.from_pretrained("BAAI/AltDiffusion"							,			subfolder="scheduler"							)

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

        UpperCAmelCase__						:    str							=							"A painting of a squirrel eating a burger"
        UpperCAmelCase__						:    Any							=							torch.manual_seed(0							)

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

        UpperCAmelCase__						:    Any							=							image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        UpperCAmelCase__						:    Union[str, Any]							=							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"""



from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCAmelCase	:	Optional[int]				=      logging.get_logger(__name__)

_lowerCAmelCase	:	Tuple				=      {
    """uw-madison/mra-base-512-4""": """https://huggingface.co/uw-madison/mra-base-512-4/resolve/main/config.json""",
}




class 			lowerCAmelCase__			(    __magic_name__   ):
   SCREAMING_SNAKE_CASE_    ='''mra'''
   def __init__(		self      :	Any							,			snake_case__      :	List[str]=5_0_2_6_5							,			snake_case__      :	Any=7_6_8							,			snake_case__      :	Union[str, Any]=1_2							,			snake_case__      :	Optional[Any]=1_2							,			snake_case__      :	Tuple=3_0_7_2							,			snake_case__      :	str="gelu"							,			snake_case__      :	Any=0.1							,			snake_case__      :	Tuple=0.1							,			snake_case__      :	Tuple=5_1_2							,			snake_case__      :	Union[str, Any]=1							,			snake_case__      :	List[Any]=0.02							,			snake_case__      :	str=1e-5							,			snake_case__      :	List[Any]="absolute"							,			snake_case__      :	str=4							,			snake_case__      :	List[str]="full"							,			snake_case__      :	Tuple=0							,			snake_case__      :	Any=0							,			snake_case__      :	Union[str, Any]=1							,			snake_case__      :	int=0							,			snake_case__      :	int=2							,			**snake_case__      :	List[Any]							,			):


        '''simple docstring'''
        super().__init__(pad_token_id=snake_case__							,			bos_token_id=snake_case__							,			eos_token_id=snake_case__							,			**snake_case__							)

        UpperCAmelCase__						:    List[Any]							=							vocab_size
        UpperCAmelCase__						:    str							=							max_position_embeddings
        UpperCAmelCase__						:    Any							=							hidden_size
        UpperCAmelCase__						:    Union[str, Any]							=							num_hidden_layers
        UpperCAmelCase__						:    str							=							num_attention_heads
        UpperCAmelCase__						:    int							=							intermediate_size
        UpperCAmelCase__						:    int							=							hidden_act
        UpperCAmelCase__						:    List[str]							=							hidden_dropout_prob
        UpperCAmelCase__						:    List[str]							=							attention_probs_dropout_prob
        UpperCAmelCase__						:    Any							=							initializer_range
        UpperCAmelCase__						:    Any							=							type_vocab_size
        UpperCAmelCase__						:    Dict							=							layer_norm_eps
        UpperCAmelCase__						:    Tuple							=							position_embedding_type
        UpperCAmelCase__						:    List[str]							=							block_per_row
        UpperCAmelCase__						:    Optional[Any]							=							approx_mode
        UpperCAmelCase__						:    Any							=							initial_prior_first_n_blocks
        UpperCAmelCase__						:    List[Any]							=							initial_prior_diagonal_n_blocks


