




import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
		    SpeechaTextaDecoder,
		    SpeechaTextaForCausalLM,
		)

@require_torch
class 	__lowerCAmelCase			:
							def __init__(		self		:List[Any]	,   __magic_name__		:Optional[int]	,   __magic_name__		:int=99	,   __magic_name__		:Optional[int]=13	,   __magic_name__		:str=16	,   __magic_name__		:List[str]=7	,   __magic_name__		:Tuple=True	,   __magic_name__		:Union[str, Any]=True	,   __magic_name__		:Optional[Any]=True	,   __magic_name__		:Tuple=False	,   __magic_name__		:Optional[int]=True	,   __magic_name__		:Dict=2	,   __magic_name__		:Dict=32	,   __magic_name__		:List[str]=4	,   __magic_name__		:Tuple=4	,   __magic_name__		:Any=30	,   __magic_name__		:int=0	,   __magic_name__		:List[str]=1	,   __magic_name__		:List[Any]=2	,   __magic_name__		:str=None	,   ):







								'''simple docstring'''







								a		       =							parent
								a		       =							batch_size
								a		       =							decoder_seq_length
								# For common tests
								a		       =							self.decoder_seq_length
								a		       =							is_training
								a		       =							use_attention_mask
								a		       =							use_labels

								a		       =							vocab_size
								a		       =							d_model
								a		       =							d_model
								a		       =							decoder_layers
								a		       =							decoder_layers
								a		       =							decoder_ffn_dim
								a		       =							decoder_attention_heads
								a		       =							decoder_attention_heads
								a		       =							eos_token_id
								a		       =							bos_token_id
								a		       =							pad_token_id
								a		       =							decoder_start_token_id
								a		       =							use_cache
								a		       =							max_position_embeddings

								a		       =							None
								a		       =							decoder_seq_length
								a		       =							2
								a		       =							1
							def    lowerCamelCase__							(		self		:Optional[int]       ):







								'''simple docstring'''







								a		       =							ids_tensor([self.batch_size, self.decoder_seq_length]	,   self.vocab_size       )

								a		       =							None
								if self.use_attention_mask:
									a		       =							ids_tensor([self.batch_size, self.decoder_seq_length]	,   vocab_size=2       )

								a		       =							None
								if self.use_labels:
									a		       =							ids_tensor([self.batch_size, self.decoder_seq_length]	,   self.vocab_size       )

								a		       =							SpeechaTextaConfig(
								    vocab_size=self.vocab_size	,   d_model=self.d_model	,   decoder_layers=self.decoder_layers	,   decoder_ffn_dim=self.decoder_ffn_dim	,   decoder_attention_heads=self.decoder_attention_heads	,   eos_token_id=self.eos_token_id	,   bos_token_id=self.bos_token_id	,   use_cache=self.use_cache	,   pad_token_id=self.pad_token_id	,   decoder_start_token_id=self.decoder_start_token_id	,   max_position_embeddings=self.max_position_embeddings	,   )

								return (
								    config,
								    input_ids,
								    attention_mask,
								    lm_labels,
								)
							def    lowerCamelCase__							(		self		:List[str]	,   __magic_name__		:Any	,   __magic_name__		:Any	,   __magic_name__		:str	,   __magic_name__		:Tuple	,   ):







								'''simple docstring'''







								a		       =							True
								a		       =							SpeechaTextaDecoder(config=__magic_name__       ).to(__magic_name__       ).eval()
								a		       =							input_ids[:2]

								input_ids[input_ids == 0] += 1
								# first forward pass
								a		       =							model(__magic_name__	,   use_cache=__magic_name__       )
								a		       =							model(__magic_name__       )
								a		       =							model(__magic_name__	,   use_cache=__magic_name__       )

								self.parent.assertTrue(len(__magic_name__       ) == len(__magic_name__       )       )
								self.parent.assertTrue(len(__magic_name__       ) == len(__magic_name__       ) + 1       )

								a		       =							outputs["""past_key_values"""]

								# create hypothetical next token and extent to next_input_ids
								a		       =							ids_tensor((2, 1)	,   config.vocab_size - 1       ) + 1

								# append to next input_ids and
								a		       =							torch.cat([input_ids, next_tokens]	,   dim=-1       )

								print(__magic_name__       )

								a		       =							model(__magic_name__       )["""last_hidden_state"""]
								a		       =							model(__magic_name__	,   past_key_values=__magic_name__       )["""last_hidden_state"""]

								# select random slice
								a		       =							ids_tensor((1,)	,   output_from_past.shape[-1]       ).item()
								a		       =							output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
								a		       =							output_from_past[:, 0, random_slice_idx].detach()

								# test that outputs are equal for slice
								assert torch.allclose(__magic_name__	,   __magic_name__	,   atol=1E-3       )





							def    lowerCamelCase__							(		self		:List[Any]       ):







								'''simple docstring'''







								a		       =							self.prepare_config_and_inputs()
								(
								    (
								    a
								)       ,				(
								    a
								)       ,				(
								    a
								)       ,				(
								    a
								)       ,				
								)		       =							config_and_inputs

								a		       =							{
								    """input_ids""": input_ids,
								    """attention_mask""": attention_mask,
								}
								return config, inputs_dict







@require_torch
class 	__lowerCAmelCase			(			__magic_name__   , __magic_name__   , __magic_name__   , unittest.TestCase				):
							UpperCamelCase__        =						(SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
							UpperCamelCase__        =						(SpeechaTextaForCausalLM,) if is_torch_available() else ()
							UpperCamelCase__        =						{'''text-generation''': SpeechaTextaForCausalLM} if is_torch_available() else {}
							UpperCamelCase__        =						True
							UpperCamelCase__        =						False
							def    lowerCamelCase__							(		self		:Optional[Any]	,   ):







								'''simple docstring'''







								a		       =							SpeechaTextaStandaloneDecoderModelTester(self	,   is_training=__magic_name__       )
								a		       =							ConfigTester(self	,   config_class=__magic_name__       )
							def    lowerCamelCase__							(		self		:Optional[int]       ):







								'''simple docstring'''







								pass
							def    lowerCamelCase__							(		self		:List[str]       ):







								'''simple docstring'''







								pass
							def    lowerCamelCase__							(		self		:Any       ):







								'''simple docstring'''







								pass
							def    lowerCamelCase__							(		self		:List[Any]       ):







								'''simple docstring'''







								self.config_tester.run_common_tests()
							def    lowerCamelCase__							(		self		:Optional[Any]       ):







								'''simple docstring'''







								a		       =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_decoder_model_past(*__magic_name__       )





							def    lowerCamelCase__							(		self		:List[str]       ):







								'''simple docstring'''







								return







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 	__lowerCAmelCase			(			__magic_name__   , __magic_name__   , 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    lowerCamelCase__							(		self		:Optional[Any]       ):







								'''simple docstring'''







								torch.manual_seed(0       )
								a		       =							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=__magic_name__	,   )
								a		       =							DDIMParallelScheduler(
								    beta_start=0.00085	,   beta_end=0.012	,   beta_schedule="""scaled_linear"""	,   clip_sample=__magic_name__	,   set_alpha_to_one=__magic_name__	,   )
								torch.manual_seed(0       )
								a		       =							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       )
								a		       =							CLIPTextConfig(
								    bos_token_id=0	,   eos_token_id=2	,   hidden_size=32	,   intermediate_size=37	,   layer_norm_eps=1E-0_5	,   num_attention_heads=4	,   num_hidden_layers=5	,   pad_token_id=1	,   vocab_size=1000	,   hidden_act="""gelu"""	,   projection_dim=512	,   )
								a		       =							CLIPTextModel(__magic_name__       )
								a		       =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

								a		       =							{
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """vae""": vae,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """safety_checker""": None,
								    """feature_extractor""": None,
								}
								return components
							def    lowerCamelCase__							(		self		:Any	,   __magic_name__		:List[Any]	,   __magic_name__		:Tuple=0       ):







								'''simple docstring'''







								if str(__magic_name__       ).startswith("""mps"""       ):
									a		       =							torch.manual_seed(__magic_name__       )
								else:
									a		       =							torch.Generator(device=__magic_name__       ).manual_seed(__magic_name__       )
								a		       =							{
								    """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    lowerCamelCase__							(		self		:Dict       ):







								'''simple docstring'''







								a		       =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
								a		       =							self.get_dummy_components()
								a		       =							StableDiffusionParadigmsPipeline(**__magic_name__       )
								a		       =							sd_pipe.to(__magic_name__       )
								sd_pipe.set_progress_bar_config(disable=__magic_name__       )

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

								a		       =							np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]       )

								assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
							def    lowerCamelCase__							(		self		:Any       ):







								'''simple docstring'''







								a		       =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
								a		       =							self.get_dummy_components()
								torch.manual_seed(0       )
								a		       =							DDPMParallelScheduler()
								torch.manual_seed(0       )
								a		       =							StableDiffusionParadigmsPipeline(**__magic_name__       )
								a		       =							sd_pipe.to(__magic_name__       )
								sd_pipe.set_progress_bar_config(disable=__magic_name__       )

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

								a		       =							np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]       )

								assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
							def    lowerCamelCase__							(		self		:Optional[int]       ):







								'''simple docstring'''







								super().test_inference_batch_consistent(batch_sizes=[1, 2]       )
							def    lowerCamelCase__							(		self		:str       ):







								'''simple docstring'''







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





							def    lowerCamelCase__							(		self		:Optional[Any]       ):







								'''simple docstring'''







								a		       =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
								a		       =							self.get_dummy_components()
								a		       =							StableDiffusionParadigmsPipeline(**__magic_name__       )
								a		       =							sd_pipe.to(__magic_name__       )
								sd_pipe.set_progress_bar_config(disable=__magic_name__       )

								a		       =							self.get_dummy_inputs(__magic_name__       )
								a		       =							"""french fries"""
								a		       =							sd_pipe(**__magic_name__	,   negative_prompt=__magic_name__       )
								a		       =							output.images
								a		       =							image[0, -3:, -3:, -1]

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

								a		       =							np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]       )

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







@slow
@require_torch_gpu
class 	__lowerCAmelCase			(			unittest.TestCase				):
							def    lowerCamelCase__							(		self		:Optional[int]       ):







								'''simple docstring'''







								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()
							def    lowerCamelCase__							(		self		:str	,   __magic_name__		:Union[str, Any]=0       ):







								'''simple docstring'''







								a		       =							torch.Generator(device=__magic_name__       ).manual_seed(__magic_name__       )
								a		       =							{
								    """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    lowerCamelCase__							(		self		:Dict       ):







								'''simple docstring'''







								a		       =							"""stabilityai/stable-diffusion-2-base"""
								a		       =							DDIMParallelScheduler.from_pretrained(__magic_name__	,   subfolder="""scheduler"""       )
								a		       =							StableDiffusionParadigmsPipeline.from_pretrained(__magic_name__	,   scheduler=__magic_name__	,   safety_checker=__magic_name__       )
								pipe.to(__magic_name__       )
								pipe.set_progress_bar_config(disable=__magic_name__       )
								pipe.enable_attention_slicing()

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

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

								a		       =							np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]       )

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

