


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class   A__							(						UpperCamelCase__				,       UpperCamelCase__				,       unittest.TestCase					):
						UpperCAmelCase     =  StableDiffusionParadigmsPipeline
						UpperCAmelCase     =  TEXT_TO_IMAGE_PARAMS
						UpperCAmelCase     =  TEXT_TO_IMAGE_BATCH_PARAMS
						UpperCAmelCase     =  TEXT_TO_IMAGE_IMAGE_PARAMS
						UpperCAmelCase     =  TEXT_TO_IMAGE_IMAGE_PARAMS




						def 							__UpperCamelCase			(    self   : str					)							->       int:


										"""simple docstring"""





										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=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			,      attention_head_dim=(2, 4)			,      use_linear_projection=_a			,      )
										_SCREAMING_SNAKE_CASE     					=DDIMParallelScheduler(
										    beta_start=0.0_00_85			,      beta_end=0.0_12			,      beta_schedule='''scaled_linear'''			,      clip_sample=_a			,      set_alpha_to_one=_a			,      )
										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=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			,      sample_size=128			,      )
										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=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			,      hidden_act='''gelu'''			,      projection_dim=512			,      )
										_SCREAMING_SNAKE_CASE     					=CLIPTextModel(_a					)
										_SCREAMING_SNAKE_CASE     					=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

										_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   : Optional[int]			,      _a   : Tuple			,      _a   : int=0					)							->       Optional[Any]:


										"""simple docstring"""





										if str(_a					).startswith('''mps'''					):
														_SCREAMING_SNAKE_CASE     					=torch.manual_seed(_a					)
										else:
														_SCREAMING_SNAKE_CASE     					=torch.Generator(device=_a					).manual_seed(_a					)
										_SCREAMING_SNAKE_CASE     					={
										    '''prompt''': '''a photograph of an astronaut riding a horse''',
										    '''generator''': generator,
										    '''num_inference_steps''': 10,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										    '''parallel''': 3,
										    '''debug''': True,
										}
										return inputs




						def 							__UpperCamelCase			(    self   : Any					)							->       Tuple:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					='''cpu'''  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE     					=self.get_dummy_components()
										_SCREAMING_SNAKE_CASE     					=StableDiffusionParadigmsPipeline(**_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe.to(_a					)
										sd_pipe.set_progress_bar_config(disable=_a					)

										_SCREAMING_SNAKE_CASE     					=self.get_dummy_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe(**_a					).images
										_SCREAMING_SNAKE_CASE     					=image[0, -3:, -3:, -1]
										assert image.shape == (1, 64, 64, 3)

										_SCREAMING_SNAKE_CASE     					=np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23]					)

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




						def 							__UpperCamelCase			(    self   : Any					)							->       Optional[Any]:


										"""simple docstring"""





										_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     					=DDPMParallelScheduler()
										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=StableDiffusionParadigmsPipeline(**_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe.to(_a					)
										sd_pipe.set_progress_bar_config(disable=_a					)

										_SCREAMING_SNAKE_CASE     					=self.get_dummy_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe(**_a					).images
										_SCREAMING_SNAKE_CASE     					=image[0, -3:, -3:, -1]
										assert image.shape == (1, 64, 64, 3)

										_SCREAMING_SNAKE_CASE     					=np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68]					)

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




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


										"""simple docstring"""





										super().test_inference_batch_consistent(batch_sizes=[1, 2]					)




						def 							__UpperCamelCase			(    self   : List[str]					)							->       Dict:


										"""simple docstring"""





										super().test_inference_batch_single_identical(batch_size=2			,      expected_max_diff=3E-3					)




						def 							__UpperCamelCase			(    self   : Optional[int]					)							->       Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					='''cpu'''  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE     					=self.get_dummy_components()
										_SCREAMING_SNAKE_CASE     					=StableDiffusionParadigmsPipeline(**_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe.to(_a					)
										sd_pipe.set_progress_bar_config(disable=_a					)

										_SCREAMING_SNAKE_CASE     					=self.get_dummy_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					='''french fries'''
										_SCREAMING_SNAKE_CASE     					=sd_pipe(**_a			,      negative_prompt=_a					)
										_SCREAMING_SNAKE_CASE     					=output.images
										_SCREAMING_SNAKE_CASE     					=image[0, -3:, -3:, -1]

										assert image.shape == (1, 64, 64, 3)

										_SCREAMING_SNAKE_CASE     					=np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15]					)

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

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




						def 							__UpperCamelCase			(    self   : str					)							->       Dict:


										"""simple docstring"""





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




						def 							__UpperCamelCase			(    self   : int			,      _a   : Optional[int]=0					)							->       Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=torch.Generator(device=_a					).manual_seed(_a					)
										_SCREAMING_SNAKE_CASE     					={
										    '''prompt''': '''a photograph of an astronaut riding a horse''',
										    '''generator''': generator,
										    '''num_inference_steps''': 10,
										    '''guidance_scale''': 7.5,
										    '''output_type''': '''numpy''',
										    '''parallel''': 3,
										    '''debug''': True,
										}
										return inputs




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					='''stabilityai/stable-diffusion-2-base'''
										_SCREAMING_SNAKE_CASE     					=DDIMParallelScheduler.from_pretrained(_a			,      subfolder='''scheduler'''					)
										_SCREAMING_SNAKE_CASE     					=StableDiffusionParadigmsPipeline.from_pretrained(_a			,      scheduler=_a			,      safety_checker=_a					)
										pipe.to(_a					)
										pipe.set_progress_bar_config(disable=_a					)
										pipe.enable_attention_slicing()

										_SCREAMING_SNAKE_CASE     					=self.get_inputs()
										_SCREAMING_SNAKE_CASE     					=pipe(**_a					).images
										_SCREAMING_SNAKE_CASE     					=image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 512, 512, 3)

										_SCREAMING_SNAKE_CASE     					=np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41]					)

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



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


snake_case_  :     Tuple							  =							logging.get_logger(__name__)

snake_case_  :     int							  =							{
    '''uclanlp/plbart-base''': '''https://huggingface.co/uclanlp/plbart-base/resolve/main/config.json''',
    # See all PLBART models at https://huggingface.co/models?filter=plbart
}

class   A__							(						UpperCamelCase__					):
						UpperCAmelCase     =  "plbart"
						UpperCAmelCase     =  ["past_key_values"]
						UpperCAmelCase     =  {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




						def __init__(    self   : Union[str, Any]			,      _a   : List[str]=5_0005			,      _a   : List[str]=1024			,      _a   : int=6			,      _a   : Union[str, Any]=3072			,      _a   : List[str]=12			,      _a   : int=6			,      _a   : Optional[int]=3072			,      _a   : Any=12			,      _a   : Any=0.0			,      _a   : Optional[Any]=0.0			,      _a   : List[Any]=True			,      _a   : Any=True			,      _a   : Any="gelu"			,      _a   : int=768			,      _a   : int=0.1			,      _a   : int=0.1			,      _a   : List[str]=0.0			,      _a   : Optional[int]=0.02			,      _a   : Dict=0.0			,      _a   : List[str]=True			,      _a   : List[Any]=1			,      _a   : Tuple=0			,      _a   : Tuple=2			,      _a   : Optional[int]=2			,      **_a   : Optional[int]			,      )							->       Optional[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=vocab_size
										_SCREAMING_SNAKE_CASE     					=max_position_embeddings
										_SCREAMING_SNAKE_CASE     					=d_model
										_SCREAMING_SNAKE_CASE     					=encoder_ffn_dim
										_SCREAMING_SNAKE_CASE     					=encoder_layers
										_SCREAMING_SNAKE_CASE     					=encoder_attention_heads
										_SCREAMING_SNAKE_CASE     					=decoder_ffn_dim
										_SCREAMING_SNAKE_CASE     					=decoder_layers
										_SCREAMING_SNAKE_CASE     					=decoder_attention_heads
										_SCREAMING_SNAKE_CASE     					=dropout
										_SCREAMING_SNAKE_CASE     					=attention_dropout
										_SCREAMING_SNAKE_CASE     					=activation_dropout
										_SCREAMING_SNAKE_CASE     					=activation_function
										_SCREAMING_SNAKE_CASE     					=init_std
										_SCREAMING_SNAKE_CASE     					=encoder_layerdrop
										_SCREAMING_SNAKE_CASE     					=decoder_layerdrop
										_SCREAMING_SNAKE_CASE     					=classifier_dropout
										_SCREAMING_SNAKE_CASE     					=use_cache
										_SCREAMING_SNAKE_CASE     					=encoder_layers
										_SCREAMING_SNAKE_CASE     					=scale_embedding  # scale factor will be sqrt(d_model) if True
										super().__init__(
										    pad_token_id=_a			,      bos_token_id=_a			,      eos_token_id=_a			,      is_encoder_decoder=_a			,      forced_eos_token_id=_a			,      **_a			,      )

class   A__							(						UpperCamelCase__					):




						@property
						def 							__UpperCamelCase			(    self   : int					)							->       Mapping[str, Mapping[int, str]]:


										"""simple docstring"""





										return OrderedDict(
										    [
										        ('''input_ids''', {0: '''batch''', 1: '''sequence'''}),
										        ('''attention_mask''', {0: '''batch''', 1: '''sequence'''}),
										    ]					)




						@property
						def 							__UpperCamelCase			(    self   : List[str]					)							->       Mapping[str, Mapping[int, str]]:


										"""simple docstring"""





										if self.use_past:
														return OrderedDict(
														    [
														        ('''last_hidden_state''', {0: '''batch''', 1: '''sequence'''}),
														        ('''past_keys''', {0: '''batch''', 2: '''sequence'''}),
														        ('''encoder_last_hidden_state''', {0: '''batch''', 1: '''sequence'''}),
														    ]					)
										else:
														return OrderedDict(
														    [
														        ('''last_hidden_state''', {0: '''batch''', 1: '''sequence'''}),
														        ('''encoder_last_hidden_state''', {0: '''batch''', 1: '''sequence'''}),
														    ]					)