


import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


_lowercase       :      Optional[Any]	=	logging.get_logger(__name__)  # pylint: disable=invalid-name






class       _UpperCamelCase    (					UpperCAmelCase__				):

				"""simple docstring"""







				def __init__(						self				,   a__				,   a__				,   a__				,   a__				,   a__				,   a__				,   a__				,   )	->				str:
					super().__init__()
					self.register_modules(
					    vae=a__				,   text_encoder=a__				,   tokenizer=a__				,   unet=a__				,   scheduler=a__				,   safety_checker=a__				,   feature_extractor=a__				,   )

				def      _UpperCAmelCase					(						self				,   a__ = "auto"   )	->				Tuple:




					if slice_size == "auto":
						# half the attention head size is usually a good trade-off between
						# speed and memory
						A								=				self.unet.config.attention_head_dim // 2
					self.unet.set_attention_slice(a__   )

				def      _UpperCAmelCase					(						self   )	->				List[str]:
					self.enable_attention_slicing(a__   )

				@torch.no_grad()
				def __call__(						self				,   a__				,   a__ = 512				,   a__ = 512				,   a__ = 50				,   a__ = 7.5				,   a__ = None				,   a__ = 1				,   a__ = 0.0				,   a__ = None				,   a__ = None				,   a__ = "pil"				,   a__ = True				,   a__ = None				,   a__ = 1				,   a__ = None				,   **a__				,   )	->				Optional[Any]:




					if isinstance(a__				,   a__   ):
						A								=				1
					elif isinstance(a__				,   a__   ):
						A								=				len(a__   )
					else:
						raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(a__   )}'   )

					if height % 8 != 0 or width % 8 != 0:
						raise ValueError(f'`height` and `width` have to be divisible by 8 but are {height} and {width}.'   )

					if (callback_steps is None) or (
					    callback_steps is not None and (not isinstance(a__				,   a__   ) or callback_steps <= 0)
					):
						raise ValueError(
						    f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
						    f' {type(a__   )}.'   )

					# get prompt text embeddings
					A								=				self.tokenizer(
					    a__				,   padding="""max_length"""				,   max_length=self.tokenizer.model_max_length				,   return_tensors="""pt"""				,   )
					A								=				text_inputs.input_ids

					if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
						A								=				self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]   )
						logger.warning(
						    """The following part of your input was truncated because CLIP can only handle sequences up to"""
						    f' {self.tokenizer.model_max_length} tokens: {removed_text}'   )
						A								=				text_input_ids[:, : self.tokenizer.model_max_length]

					if text_embeddings is None:
						A								=				self.text_encoder(text_input_ids.to(self.device   )   )[0]

					# duplicate text embeddings for each generation per prompt, using mps friendly method
					A								=				text_embeddings.shape
					A								=				text_embeddings.repeat(1				,   a__				,   1   )
					A								=				text_embeddings.view(bs_embed * num_images_per_prompt				,   a__				,   -1   )

					# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
					# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
					# corresponds to doing no classifier free guidance.
					A								=				guidance_scale > 1.0
					# get unconditional embeddings for classifier free guidance
					if do_classifier_free_guidance:
						A								=				42
						if negative_prompt is None:
							A								=				[""""""]
						elif type(a__   ) is not type(a__   ):
							raise TypeError(
							    f'`negative_prompt` should be the same type to `prompt`, but got {type(a__   )} !='
							    f' {type(a__   )}.'   )
						elif isinstance(a__				,   a__   ):
							A								=				[negative_prompt]
						elif batch_size != len(a__   ):
							raise ValueError(
							    f'`negative_prompt`: {negative_prompt} has batch size {len(a__   )}, but `prompt`:'
							    f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
							    """ the batch size of `prompt`."""   )
						else:
							A								=				negative_prompt

						A								=				text_input_ids.shape[-1]
						A								=				self.tokenizer(
						    a__				,   padding="""max_length"""				,   max_length=a__				,   truncation=a__				,   return_tensors="""pt"""				,   )
						A								=				self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]

						# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
						A								=				uncond_embeddings.shape[1]
						A								=				uncond_embeddings.repeat(a__				,   a__				,   1   )
						A								=				uncond_embeddings.view(batch_size * num_images_per_prompt				,   a__				,   -1   )

						# For classifier free guidance, we need to do two forward passes.
						# Here we concatenate the unconditional and text embeddings into a single batch
						# to avoid doing two forward passes
						A								=				torch.cat([uncond_embeddings, text_embeddings]   )

					# get the initial random noise unless the user supplied it

					# Unlike in other pipelines, latents need to be generated in the target device
					# for 1-to-1 results reproducibility with the CompVis implementation.
					# However this currently doesn't work in `mps`.
					A								=				(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
					A								=				(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
					A								=				text_embeddings.dtype
					if latents is None:
						if self.device.type == "mps":
							# randn does not exist on mps
							A								=				torch.randn(
							    a__				,   generator=a__				,   device="""cpu"""				,   dtype=a__   ).to(self.device   )
							A								=				torch.randn(a__				,   generator=a__				,   device="""cpu"""				,   dtype=a__   ).to(
							    self.device   )
						else:
							A								=				torch.randn(
							    a__				,   generator=a__				,   device=self.device				,   dtype=a__   )
							A								=				torch.randn(a__				,   generator=a__				,   device=self.device				,   dtype=a__   )
					else:
						if latents_reference.shape != latents_shape:
							raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'   )
						A								=				latents_reference.to(self.device   )
						A								=				latents.to(self.device   )

					# This is the key part of the pipeline where we
					# try to ensure that the generated images w/ the same seed
					# but different sizes actually result in similar images
					A								=				(latents_shape[3] - latents_shape_reference[3]) // 2
					A								=				(latents_shape[2] - latents_shape_reference[2]) // 2
					A								=				latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
					A								=				latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
					A								=				0 if dx < 0 else dx
					A								=				0 if dy < 0 else dy
					A								=				max(-dx				,   0   )
					A								=				max(-dy				,   0   )
					# import pdb
					# pdb.set_trace()
					A								=				latents_reference[:, :, dy : dy + h, dx : dx + w]

					# set timesteps
					self.scheduler.set_timesteps(a__   )

					# Some schedulers like PNDM have timesteps as arrays
					# It's more optimized to move all timesteps to correct device beforehand
					A								=				self.scheduler.timesteps.to(self.device   )

					# scale the initial noise by the standard deviation required by the scheduler
					A								=				latents * self.scheduler.init_noise_sigma

					# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
					# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
					# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
					# and should be between [0, 1]
					A								=				"""eta""" in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
					A								=				{}
					if accepts_eta:
						A								=				eta

					for i, t in enumerate(self.progress_bar(a__   )   ):
						# expand the latents if we are doing classifier free guidance
						A								=				torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents
						A								=				self.scheduler.scale_model_input(a__				,   a__   )

						# predict the noise residual
						A								=				self.unet(a__				,   a__				,   encoder_hidden_states=a__   ).sample

						# perform guidance
						if do_classifier_free_guidance:
							A								=				noise_pred.chunk(2   )
							A								=				noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

						# compute the previous noisy sample x_t -> x_t-1
						A								=				self.scheduler.step(a__				,   a__				,   a__				,   **a__   ).prev_sample

						# call the callback, if provided
						if callback is not None and i % callback_steps == 0:
							callback(a__				,   a__				,   a__   )

					A								=				1 / 0.1_82_15 * latents
					A								=				self.vae.decode(a__   ).sample

					A								=				(image / 2 + 0.5).clamp(0				,   1   )

					# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
					A								=				image.cpu().permute(0				,   2				,   3				,   1   ).float().numpy()

					if self.safety_checker is not None:
						A								=				self.feature_extractor(self.numpy_to_pil(a__   )				,   return_tensors="""pt"""   ).to(
						    self.device   )
						A								=				self.safety_checker(
						    images=a__				,   clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype   )   )
					else:
						A								=				None

					if output_type == "pil":
						A								=				self.numpy_to_pil(a__   )

					if not return_dict:
						return (image, has_nsfw_concept)

					return StableDiffusionPipelineOutput(images=a__				,   nsfw_content_detected=a__   )







import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
				from transformers import SpeechaTextFeatureExtractor

_lowercase       :      Tuple	=	random.Random()





def   _lowerCAmelCase				(   UpperCamelCase__:     List[str]							,					UpperCamelCase__:     str=1.0							,					UpperCamelCase__:     int=None							,					UpperCamelCase__:     Optional[Any]=None						) ->   str:





	"""simple docstring"""







	if rng is None:
		A								=				global_rng

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

	return values






@require_torch
@require_torchaudio
class       _UpperCamelCase    (					unittest.TestCase				):

				"""simple docstring"""







				def __init__(						self				,   a__				,   a__=7				,   a__=400				,   a__=2000				,   a__=24				,   a__=24				,   a__=0.0				,   a__=1_6000				,   a__=True				,   a__=True				,   )	->				Any:
					A								=				parent
					A								=				batch_size
					A								=				min_seq_length
					A								=				max_seq_length
					A								=				(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
					A								=				feature_size
					A								=				num_mel_bins
					A								=				padding_value
					A								=				sampling_rate
					A								=				return_attention_mask
					A								=				do_normalize

				def      _UpperCAmelCase					(						self   )	->				Optional[int]:
					return {
					    "feature_size": self.feature_size,
					    "num_mel_bins": self.num_mel_bins,
					    "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   )	->				Tuple:
					def _flatten(a__   ):
						return list(itertools.chain(*a__   )   )

					if equal_length:
						A								=				[floats_list((self.max_seq_length, self.feature_size)   ) for _ in range(self.batch_size   )]
					else:
						# make sure that inputs increase in size
						A								=				[
						    floats_list((x, self.feature_size)   )
						    for x in range(self.min_seq_length				,   self.max_seq_length				,   self.seq_length_diff   )
						]
					if numpify:
						A								=				[np.asarray(a__   ) for x in speech_inputs]
					return speech_inputs






@require_torch
@require_torchaudio
class       _UpperCamelCase    (					__snake_case    ,		unittest.TestCase				):

				"""simple docstring"""







				lowerCAmelCase 		= SpeechaTextFeatureExtractor if is_speech_available() else None

				def      _UpperCAmelCase					(						self   )	->				Dict:
					A								=				SpeechaTextFeatureExtractionTester(self   )

				def      _UpperCAmelCase					(						self				,   a__   )	->				Any:
					self.assertTrue(np.all(np.mean(a__				,   axis=0   ) < 1e-3   )   )
					self.assertTrue(np.all(np.abs(np.var(a__				,   axis=0   ) - 1   ) < 1e-3   )   )

				def      _UpperCAmelCase					(						self   )	->				Dict:
					# Tests that all call wrap to encode_plus and batch_encode_plus
					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					# create three inputs of length 800, 1000, and 1200
					A								=				[floats_list((1, x)   )[0] for x in range(800				,   1400				,   200   )]
					A								=				[np.asarray(a__   ) for speech_input in speech_inputs]

					# Test feature size
					A								=				feature_extractor(a__				,   padding=a__				,   return_tensors="""np"""   ).input_features
					self.assertTrue(input_features.ndim == 3   )
					self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size   )

					# Test not batched input
					A								=				feature_extractor(speech_inputs[0]				,   return_tensors="""np"""   ).input_features
					A								=				feature_extractor(np_speech_inputs[0]				,   return_tensors="""np"""   ).input_features
					self.assertTrue(np.allclose(a__				,   a__				,   atol=1e-3   )   )

					# Test batched
					A								=				feature_extractor(a__				,   return_tensors="""np"""   ).input_features
					A								=				feature_extractor(a__				,   return_tensors="""np"""   ).input_features
					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.
					A								=				[floats_list((1, x)   )[0] for x in (800, 800, 800)]
					A								=				np.asarray(a__   )
					A								=				feature_extractor(a__				,   return_tensors="""np"""   ).input_features
					A								=				feature_extractor(a__				,   return_tensors="""np"""   ).input_features
					for enc_seq_a, enc_seq_a in zip(a__				,   a__   ):
						self.assertTrue(np.allclose(a__				,   a__				,   atol=1e-3   )   )

				def      _UpperCAmelCase					(						self   )	->				Tuple:
					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					A								=				[floats_list((1, x)   )[0] for x in range(800				,   1400				,   200   )]

					A								=				["""longest""", """max_length""", """do_not_pad"""]
					A								=				[None, 16, None]
					for max_length, padding in zip(a__				,   a__   ):
						A								=				feature_extractor(
						    a__				,   padding=a__				,   max_length=a__				,   return_attention_mask=a__   )
						A								=				inputs.input_features
						A								=				inputs.attention_mask
						A								=				[np.sum(a__   ) for x in attention_mask]

						self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]   )
						self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]   )
						self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]   )

				def      _UpperCAmelCase					(						self   )	->				Any:
					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					A								=				[floats_list((1, x)   )[0] for x in range(800				,   1400				,   200   )]

					A								=				["""longest""", """max_length""", """do_not_pad"""]
					A								=				[None, 16, None]
					for max_length, padding in zip(a__				,   a__   ):
						A								=				feature_extractor(
						    a__				,   max_length=a__				,   padding=a__				,   return_tensors="""np"""				,   return_attention_mask=a__   )
						A								=				inputs.input_features
						A								=				inputs.attention_mask
						A								=				[np.sum(a__   ) for x in attention_mask]

						self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]   )
						self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1e-6   )
						self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]   )
						self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1e-6   )
						self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]   )

				def      _UpperCAmelCase					(						self   )	->				str:
					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					A								=				[floats_list((1, x)   )[0] for x in range(800				,   1400				,   200   )]
					A								=				feature_extractor(
					    a__				,   padding="""max_length"""				,   max_length=4				,   truncation=a__				,   return_tensors="""np"""				,   return_attention_mask=a__				,   )
					A								=				inputs.input_features
					A								=				inputs.attention_mask
					A								=				np.sum(attention_mask == 1				,   axis=1   )

					self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
					self._check_zero_mean_unit_variance(input_features[1]   )
					self._check_zero_mean_unit_variance(input_features[2]   )

				def      _UpperCAmelCase					(						self   )	->				Tuple:
					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					A								=				[floats_list((1, x)   )[0] for x in range(800				,   1400				,   200   )]
					A								=				feature_extractor(
					    a__				,   padding="""longest"""				,   max_length=4				,   truncation=a__				,   return_tensors="""np"""				,   return_attention_mask=a__				,   )
					A								=				inputs.input_features
					A								=				inputs.attention_mask
					A								=				np.sum(attention_mask == 1				,   axis=1   )

					self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
					self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]   )
					self._check_zero_mean_unit_variance(input_features[2]   )

					# make sure that if max_length < longest -> then pad to max_length
					self.assertEqual(input_features.shape				,   (3, 4, 24)   )

					A								=				[floats_list((1, x)   )[0] for x in range(800				,   1400				,   200   )]
					A								=				feature_extractor(
					    a__				,   padding="""longest"""				,   max_length=16				,   truncation=a__				,   return_tensors="""np"""				,   return_attention_mask=a__				,   )
					A								=				inputs.input_features
					A								=				inputs.attention_mask
					A								=				np.sum(attention_mask == 1				,   axis=1   )

					self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
					self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]   )
					self._check_zero_mean_unit_variance(input_features[2]   )

					# make sure that if max_length < longest -> then pad to max_length
					self.assertEqual(input_features.shape				,   (3, 6, 24)   )

				def      _UpperCAmelCase					(						self   )	->				Any:
					import torch

					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					A								=				np.random.rand(100				,   32   ).astype(np.floataa   )
					A								=				np_speech_inputs.tolist()

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

				def      _UpperCAmelCase					(						self				,   a__   )	->				Optional[int]:
					from datasets import load_dataset

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

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

				def      _UpperCAmelCase					(						self   )	->				Any:
					# fmt: off
					A								=				np.array([
					    -1.57_45, -1.77_13, -1.70_20, -1.60_69, -1.22_50, -1.11_05, -0.90_72, -0.82_41,
					    -1.23_10, -0.80_98, -0.33_20, -0.41_01, -0.79_85, -0.49_96, -0.82_13, -0.91_28,
					    -1.04_20, -1.12_86, -1.04_40, -0.79_99, -0.84_05, -1.22_75, -1.54_43, -1.46_25,
					]   )
					# fmt: on

					A								=				self._load_datasamples(1   )
					A								=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
					A								=				feature_extractor(a__				,   return_tensors="""pt"""   ).input_features
					self.assertEquals(input_features.shape				,   (1, 584, 24)   )
					self.assertTrue(np.allclose(input_features[0, 0, :30]				,   a__				,   atol=1e-4   )   )



