






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       __SCREAMING_SNAKE_CASE   (				__lowercase						,   __lowercase						,   __lowercase						,   unittest.TestCase):
						_SCREAMING_SNAKE_CASE		:    Optional[int]						=   AltDiffusionPipeline
						_SCREAMING_SNAKE_CASE		:    str						=   TEXT_TO_IMAGE_PARAMS
						_SCREAMING_SNAKE_CASE		:    Union[str, Any]						=   TEXT_TO_IMAGE_BATCH_PARAMS
						_SCREAMING_SNAKE_CASE		:    Union[str, Any]						=   TEXT_TO_IMAGE_IMAGE_PARAMS
						_SCREAMING_SNAKE_CASE		:    int						=   TEXT_TO_IMAGE_IMAGE_PARAMS






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  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       ,  )
												lowerCAmelCase__        =  DDIMScheduler(
												    beta_start=0.0_00_85       ,  beta_end=0.0_12       ,  beta_schedule='scaled_linear'       ,  clip_sample=_UpperCamelCase       ,  set_alpha_to_one=_UpperCamelCase       ,  )
												torch.manual_seed(0				)
												lowerCAmelCase__        =  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				)
												lowerCAmelCase__        =  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=50_02       ,  )
												lowerCAmelCase__        =  CLIPTextModel(_UpperCamelCase				)

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

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






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase=0				):






												"""simple docstring"""
												if str(_UpperCamelCase				).startswith('mps'				):
																		lowerCAmelCase__        =  torch.manual_seed(_UpperCamelCase				)
												else:
																		lowerCAmelCase__        =  torch.Generator(device=_UpperCamelCase				).manual_seed(_UpperCamelCase				)
												lowerCAmelCase__        =  {
												    '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				):






												"""simple docstring"""
												super().test_attention_slicing_forward_pass(expected_max_diff=3E-3				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												super().test_inference_batch_single_identical(expected_max_diff=3E-3				)






						def  UpperCamelCase__ (		self				):






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

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

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

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

												assert image.shape == (1, 64, 64, 3)
												lowerCAmelCase__        =  np.array(
												    [0.5_74_81_62, 0.60_44_71_45, 0.48_82_12_17, 0.50_10_06_36, 0.5_43_11_85, 0.45_76_36_83, 0.49_65_76_96, 0.48_13_27_33, 0.47_57_30_93]				)

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







						def  UpperCamelCase__ (		self				):






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

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

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

												assert image.shape == (1, 64, 64, 3)
												lowerCAmelCase__        =  np.array(
												    [0.51_60_50_93, 0.5_70_72_41, 0.47_36_55_07, 0.50_57_88_86, 0.5_63_38_77, 0.4_64_25_03, 0.5_18_20_81, 0.48_76_34_84, 0.49_08_42_37]				)

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




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






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												# make sure here that pndm scheduler skips prk
												lowerCAmelCase__        =  AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'       ,  safety_checker=_UpperCamelCase				)
												lowerCAmelCase__        =  alt_pipe.to(_UpperCamelCase				)
												alt_pipe.set_progress_bar_config(disable=_UpperCamelCase				)

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

												lowerCAmelCase__        =  output.images

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

												assert image.shape == (1, 5_12, 5_12, 3)
												lowerCAmelCase__        =  np.array([0.10_10, 0.08_00, 0.07_94, 0.08_85, 0.08_43, 0.07_62, 0.07_69, 0.07_29, 0.05_86]				)

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







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  DDIMScheduler.from_pretrained('BAAI/AltDiffusion'       ,  subfolder='scheduler'				)

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

												lowerCAmelCase__        =  'A painting of a squirrel eating a burger'
												lowerCAmelCase__        =  torch.manual_seed(0				)

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

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

												assert image.shape == (1, 5_12, 5_12, 3)
												lowerCAmelCase__        =  np.array([0.40_19, 0.40_52, 0.38_10, 0.41_19, 0.39_16, 0.39_82, 0.46_51, 0.41_95, 0.53_23]				)

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













# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from ..models.speechta import SpeechTaForTextToSpeech, SpeechTaHifiGan, SpeechTaProcessor
from ..utils import is_datasets_available
from .base import PipelineTool


if is_datasets_available():
					from datasets import load_dataset


class       __SCREAMING_SNAKE_CASE   (				__lowercase):
						_SCREAMING_SNAKE_CASE		:    List[Any]						=   '''microsoft/speecht5_tts'''
						_SCREAMING_SNAKE_CASE		:    Any						=   (
						    '''This is a tool that reads an English text out loud. It takes an input named `text` which should contain the '''
						    '''text to read (in English) and returns a waveform object containing the sound.'''
						)
						_SCREAMING_SNAKE_CASE		:    int						=   '''text_reader'''
						_SCREAMING_SNAKE_CASE		:    List[str]						=   SpeechTaProcessor
						_SCREAMING_SNAKE_CASE		:    Optional[int]						=   SpeechTaForTextToSpeech
						_SCREAMING_SNAKE_CASE		:    List[Any]						=   SpeechTaHifiGan

						_SCREAMING_SNAKE_CASE		:    Optional[int]						=   ['''text''']
						_SCREAMING_SNAKE_CASE		:    List[Any]						=   ['''audio''']






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												if self.post_processor is None:
																		lowerCAmelCase__        =  'microsoft/speecht5_hifigan'
												super().setup()






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase=None				):






												"""simple docstring"""
												lowerCAmelCase__        =  self.pre_processor(text=_UpperCamelCase       ,  return_tensors='pt'       ,  truncation=_UpperCamelCase				)

												if speaker_embeddings is None:
																		if not is_datasets_available():
																								raise ImportError('Datasets needs to be installed if not passing speaker embeddings.'				)

																		lowerCAmelCase__        =  load_dataset('Matthijs/cmu-arctic-xvectors'       ,  split='validation'				)
																		lowerCAmelCase__        =  torch.tensor(embeddings_dataset[73_05]['xvector']				).unsqueeze(0				)

												return {"input_ids": inputs["input_ids"], "speaker_embeddings": speaker_embeddings}






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												with torch.no_grad():
																		return self.model.generate_speech(**_UpperCamelCase				)







						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												with torch.no_grad():
																		return self.post_processor(_UpperCamelCase				).cpu().detach()





