"""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_       (						__snake_case			,							__snake_case			,							__snake_case			,							unittest.TestCase):
							snake_case__        =      AltDiffusionPipeline
							snake_case__        =      TEXT_TO_IMAGE_PARAMS
							snake_case__        =      TEXT_TO_IMAGE_BATCH_PARAMS
							snake_case__        =      TEXT_TO_IMAGE_IMAGE_PARAMS
							snake_case__        =      TEXT_TO_IMAGE_IMAGE_PARAMS






							def       _UpperCamelCase			(		self :      int  )       ->       List[str]:
												torch.manual_seed(0  )
												_UpperCamelCase      						=	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     ,						)
												_UpperCamelCase      						=	DDIMScheduler(
												    beta_start=0.0_0_0_8_5     ,						beta_end=0.0_1_2     ,						beta_schedule='''scaled_linear'''     ,						clip_sample=A_     ,						set_alpha_to_one=A_     ,						)
												torch.manual_seed(0  )
												_UpperCamelCase      						=	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  )
												_UpperCamelCase      						=	CLIPTextConfig(
												    bos_token_id=0     ,						eos_token_id=2     ,						hidden_size=32     ,						projection_dim=32     ,						intermediate_size=37     ,						layer_norm_eps=1E-05     ,						num_attention_heads=4     ,						num_hidden_layers=5     ,						pad_token_id=1     ,						vocab_size=5002     ,						)
												_UpperCamelCase      						=	CLIPTextModel(A_  )

												_UpperCamelCase      						=	XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''  )
												_UpperCamelCase      						=	77

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






							def       _UpperCamelCase			(		self :      Optional[Any]     ,						__UpperCamelCase :      Dict     ,						__UpperCamelCase :      Dict=0  )       ->       str:




												if str(A_  ).startswith('''mps'''  ):
																	_UpperCamelCase      						=	torch.manual_seed(A_  )
												else:
																	_UpperCamelCase      						=	torch.Generator(device=A_  ).manual_seed(A_  )
												_UpperCamelCase      						=	{
												    "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 :      Union[str, Any]  )       ->       Optional[int]:
												super().test_attention_slicing_forward_pass(expected_max_diff=3E-3  )






							def       _UpperCamelCase			(		self :      Tuple  )       ->       List[str]:
												super().test_inference_batch_single_identical(expected_max_diff=3E-3  )






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

												_UpperCamelCase      						=	self.get_dummy_components()
												torch.manual_seed(0  )
												_UpperCamelCase      						=	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     ,						)
												# TODO: remove after fixing the non-deterministic text encoder
												_UpperCamelCase      						=	RobertaSeriesModelWithTransformation(A_  )
												_UpperCamelCase      						=	text_encoder

												_UpperCamelCase      						=	AltDiffusionPipeline(**A_  )
												_UpperCamelCase      						=	alt_pipe.to(A_  )
												alt_pipe.set_progress_bar_config(disable=A_  )

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

												assert image.shape == (1, 64, 64, 3)
												_UpperCamelCase      						=	np.array(
												    [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]  )

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






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

												_UpperCamelCase      						=	self.get_dummy_components()
												_UpperCamelCase      						=	PNDMScheduler(skip_prk_steps=A_  )
												torch.manual_seed(0  )
												_UpperCamelCase      						=	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     ,						)
												# TODO: remove after fixing the non-deterministic text encoder
												_UpperCamelCase      						=	RobertaSeriesModelWithTransformation(A_  )
												_UpperCamelCase      						=	text_encoder
												_UpperCamelCase      						=	AltDiffusionPipeline(**A_  )
												_UpperCamelCase      						=	alt_pipe.to(A_  )
												alt_pipe.set_progress_bar_config(disable=A_  )

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

												assert image.shape == (1, 64, 64, 3)
												_UpperCamelCase      						=	np.array(
												    [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]  )

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







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






							def       _UpperCamelCase			(		self :      str  )       ->       List[Any]:
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






							def       _UpperCamelCase			(		self :      Optional[Any]  )       ->       Any:
												_UpperCamelCase      						=	AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''     ,						safety_checker=A_  )
												_UpperCamelCase      						=	alt_pipe.to(A_  )
												alt_pipe.set_progress_bar_config(disable=A_  )

												_UpperCamelCase      						=	"A painting of a squirrel eating a burger"
												_UpperCamelCase      						=	torch.manual_seed(0  )
												_UpperCamelCase      						=	alt_pipe([prompt]     ,						generator=A_     ,						guidance_scale=6.0     ,						num_inference_steps=20     ,						output_type='''np'''  )

												_UpperCamelCase      						=	output.images

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

												assert image.shape == (1, 512, 512, 3)
												_UpperCamelCase      						=	np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]  )

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






							def       _UpperCamelCase			(		self :      Any  )       ->       str:
												_UpperCamelCase      						=	DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''     ,						subfolder='''scheduler'''  )

												_UpperCamelCase      						=	AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''     ,						scheduler=A_     ,						safety_checker=A_  )
												_UpperCamelCase      						=	alt_pipe.to(A_  )
												alt_pipe.set_progress_bar_config(disable=A_  )

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

												_UpperCamelCase      						=	alt_pipe([prompt]     ,						generator=A_     ,						num_inference_steps=2     ,						output_type='''numpy'''  )
												_UpperCamelCase      						=	output.images

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

												assert image.shape == (1, 512, 512, 3)
												_UpperCamelCase      						=	np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]  )

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







import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


__lowerCamelCase       :     str   =     random.Random()

if is_torch_available():
    import torch



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase=1.0							,   _lowerCAmelCase=None							,   _lowerCAmelCase=None     )						-> Optional[Any]:
 if rng is None:
  UpperCamelCase	:					Optional[int]       	=   global_rng

 UpperCamelCase	:					Optional[Any]       	=   []
 for batch_idx in range(shape[0]     ):
  values.append([]     )
  for _ in range(shape[1]     ):
   values[-1].append(rng.random() * scale     )

 return values



class 					A__   (       unittest.TestCase			):





      def __init__( self       ,							A_       ,							A_=7       ,							A_=400       ,							A_=2000       ,							A_=1       ,							A_=0.0       ,							A_=1_6000       ,							A_=True       ,							A_=True       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   parent
       UpperCamelCase	:					List[Any]       	=   batch_size
       UpperCamelCase	:					List[Any]       	=   min_seq_length
       UpperCamelCase	:					List[str]       	=   max_seq_length
       UpperCamelCase	:					int       	=   (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       UpperCamelCase	:					Union[str, Any]       	=   feature_size
       UpperCamelCase	:					List[str]       	=   padding_value
       UpperCamelCase	:					Optional[Any]       	=   sampling_rate
       UpperCamelCase	:					List[str]       	=   return_attention_mask
       UpperCamelCase	:					List[Any]       	=   do_normalize





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return {
           "feature_size": self.feature_size,
           "padding_value": self.padding_value,
           "sampling_rate": self.sampling_rate,
           "return_attention_mask": self.return_attention_mask,
           "do_normalize": self.do_normalize,
       }





      def        __UpperCamelCase( self       ,							A_=False       ,							A_=False   ):
       '''simple docstring'''




       def _flatten(A_   ):
        return list(itertools.chain(*A_   )   )

       if equal_length:
        UpperCamelCase	:					List[str]       	=   floats_list((self.batch_size, self.max_seq_length)   )
       else:
        # make sure that inputs increase in size
        UpperCamelCase	:					Dict       	=   [
            _flatten(floats_list((x, self.feature_size)   )   )
            for x in range(self.min_seq_length       ,							self.max_seq_length       ,							self.seq_length_diff   )
        ]

       if numpify:
        UpperCamelCase	:					Union[str, Any]       	=   [np.asarray(A_   ) for x in speech_inputs]

       return speech_inputs



@require_torch
@require_torchaudio
class 					A__   (       __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Optional[Any]						    =   ASTFeatureExtractor





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   ASTFeatureExtractionTester(self   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
       # create three inputs of length 800, 1000, and 1200
       UpperCamelCase	:					Tuple       	=   [floats_list((1, x)   )[0] for x in range(800       ,							1400       ,							200   )]
       UpperCamelCase	:					Dict       	=   [np.asarray(A_   ) for speech_input in speech_inputs]

       # Test not batched input
       UpperCamelCase	:					Dict       	=   feat_extract(speech_inputs[0]       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					Union[str, Any]       	=   feat_extract(np_speech_inputs[0]       ,							return_tensors="np"   ).input_values
       self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )

       # Test batched
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							padding=A_       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							padding=A_       ,							return_tensors="np"   ).input_values
       for enc_seq_a, enc_seq_a in zip(A_       ,							A_   ):
        self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )

       # Test 2-D numpy arrays are batched.
       UpperCamelCase	:					Dict       	=   [floats_list((1, x)   )[0] for x in (800, 800, 800)]
       UpperCamelCase	:					int       	=   np.asarray(A_   )
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					List[str]       	=   feat_extract(A_       ,							return_tensors="np"   ).input_values
       for enc_seq_a, enc_seq_a in zip(A_       ,							A_   ):
        self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       import torch

       UpperCamelCase	:					List[Any]       	=   self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
       UpperCamelCase	:					int       	=   np.random.rand(100   ).astype(np.floataa   )
       UpperCamelCase	:					str       	=   np_speech_inputs.tolist()

       for inputs in [py_speech_inputs, np_speech_inputs]:
        UpperCamelCase	:					List[Any]       	=   feature_extractor.pad([{"input_values": inputs}]       ,							return_tensors="np"   )
        self.assertTrue(np_processed.input_values.dtype == np.floataa   )
        UpperCamelCase	:					List[str]       	=   feature_extractor.pad([{"input_values": inputs}]       ,							return_tensors="pt"   )
        self.assertTrue(pt_processed.input_values.dtype == torch.floataa   )





      def        __UpperCamelCase( self       ,							A_   ):
       '''simple docstring'''




       from datasets import load_dataset

       UpperCamelCase	:					Dict       	=   load_dataset("hf-internal-testing/librispeech_asr_dummy"       ,							"clean"       ,							split="validation"   )
       # automatic decoding with librispeech
       UpperCamelCase	:					Any       	=   ds.sort("id"   ).select(range(A_   )   )[:num_samples]["audio"]

       return [x["array"] for x in speech_samples]





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   torch.tensor(
           [-0.98_94, -1.27_76, -0.90_66, -1.27_76, -0.93_49, -1.26_09, -1.03_86, -1.27_76,
            -1.15_61, -1.27_76, -1.20_52, -1.27_23, -1.21_90, -1.21_32, -1.27_76, -1.11_33,
            -1.19_53, -1.13_43, -1.15_84, -1.22_03, -1.17_70, -1.24_74, -1.23_81, -1.19_36,
            -0.92_70, -0.83_17, -0.80_49, -0.77_06, -0.75_65, -0.78_69]   )
       # fmt: on

       UpperCamelCase	:					List[Any]       	=   self._load_datasamples(1   )
       UpperCamelCase	:					Tuple       	=   ASTFeatureExtractor()
       UpperCamelCase	:					str       	=   feature_extractor(A_       ,							return_tensors="pt"   ).input_values
       self.assertEquals(input_values.shape       ,							(1, 1024, 128)   )
       self.assertTrue(torch.allclose(input_values[0, 0, :30]       ,							A_       ,							atol=1e-4   )   )
