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

import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextModel,
    CLIPTokenizer,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import logging


_A  =	logging.get_logger(__name__)  # pylint: disable=invalid-name





class       _lowerCAmelCase   (      __a			):




				def __init__(   self							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							)							-> Tuple:
							super().__init__()

							if safety_checker is None:
										logger.warning(
										    f"""You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"""
										    " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
										    " results in services or applications open to the public. Both the diffusers team and Hugging Face"
										    " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
										    " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
										    " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."	)

							self.register_modules(
							    speech_model=_UpperCamelCase							,							speech_processor=_UpperCamelCase							,							vae=_UpperCamelCase							,							text_encoder=_UpperCamelCase							,							tokenizer=_UpperCamelCase							,							unet=_UpperCamelCase							,							scheduler=_UpperCamelCase							,							feature_extractor=_UpperCamelCase							,							)




				def 	__a		(   self							,							_UpperCamelCase = "auto"	)							-> Dict:
							if slice_size == "auto":
										lowerCAmelCase_         =      self.unet.config.attention_head_dim // 2
							self.unet.set_attention_slice(_UpperCamelCase	)




				def 	__a		(   self	)							-> List[str]:
							self.enable_attention_slicing(_UpperCamelCase	)




				@torch.no_grad()
				def __call__(   self							,							_UpperCamelCase							,							_UpperCamelCase=16_000							,							_UpperCamelCase = 512							,							_UpperCamelCase = 512							,							_UpperCamelCase = 50							,							_UpperCamelCase = 7.5							,							_UpperCamelCase = None							,							_UpperCamelCase = 1							,							_UpperCamelCase = 0.0							,							_UpperCamelCase = None							,							_UpperCamelCase = None							,							_UpperCamelCase = "pil"							,							_UpperCamelCase = True							,							_UpperCamelCase = None							,							_UpperCamelCase = 1							,							**_UpperCamelCase							,							)							-> str:
							lowerCAmelCase_         =      self.speech_processor.feature_extractor(
							    _UpperCamelCase							,							return_tensors="pt"							,							sampling_rate=_UpperCamelCase	).input_features.to(self.device	)
							lowerCAmelCase_         =      self.speech_model.generate(_UpperCamelCase							,							max_length=480_000	)

							lowerCAmelCase_         =      self.speech_processor.tokenizer.batch_decode(_UpperCamelCase							,							skip_special_tokens=_UpperCamelCase							,							normalize=_UpperCamelCase	)[
							    0
							]

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

							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(_UpperCamelCase							,							_UpperCamelCase	) or callback_steps <= 0)
							):
										raise ValueError(
										    f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
										    f""" {type(_UpperCamelCase	)}."""	)

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

							if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
										lowerCAmelCase_         =      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}"""	)
										lowerCAmelCase_         =      text_input_ids[:, : self.tokenizer.model_max_length]
							lowerCAmelCase_         =      self.text_encoder(text_input_ids.to(self.device	)	)[0]

							# duplicate text embeddings for each generation per prompt, using mps friendly method
							lowerCAmelCase_      ,    lowerCAmelCase_      ,    lowerCAmelCase_         =      text_embeddings.shape
							lowerCAmelCase_         =      text_embeddings.repeat(1							,							_UpperCamelCase							,							1	)
							lowerCAmelCase_         =      text_embeddings.view(bs_embed * num_images_per_prompt							,							_UpperCamelCase							,							-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.
							lowerCAmelCase_         =      guidance_scale > 1.0
							# get unconditional embeddings for classifier free guidance
							if do_classifier_free_guidance:
										lowerCAmelCase_         =      42
										if negative_prompt is None:
													lowerCAmelCase_         =      [""] * batch_size
										elif type(_UpperCamelCase	) is not type(_UpperCamelCase	):
													raise TypeError(
													    f"""`negative_prompt` should be the same type to `prompt`, but got {type(_UpperCamelCase	)} !="""
													    f""" {type(_UpperCamelCase	)}."""	)
										elif isinstance(_UpperCamelCase							,							_UpperCamelCase	):
													lowerCAmelCase_         =      [negative_prompt]
										elif batch_size != len(_UpperCamelCase	):
													raise ValueError(
													    f"""`negative_prompt`: {negative_prompt} has batch size {len(_UpperCamelCase	)}, but `prompt`:"""
													    f""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
													    " the batch size of `prompt`."	)
										else:
													lowerCAmelCase_         =      negative_prompt

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

										# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
										lowerCAmelCase_         =      uncond_embeddings.shape[1]
										lowerCAmelCase_         =      uncond_embeddings.repeat(1							,							_UpperCamelCase							,							1	)
										lowerCAmelCase_         =      uncond_embeddings.view(batch_size * num_images_per_prompt							,							_UpperCamelCase							,							-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
										lowerCAmelCase_         =      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`.
							lowerCAmelCase_         =      (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
							lowerCAmelCase_         =      text_embeddings.dtype
							if latents is None:
										if self.device.type == "mps":
													# randn does not exist on mps
													lowerCAmelCase_         =      torch.randn(_UpperCamelCase							,							generator=_UpperCamelCase							,							device="cpu"							,							dtype=_UpperCamelCase	).to(
													    self.device	)
										else:
													lowerCAmelCase_         =      torch.randn(_UpperCamelCase							,							generator=_UpperCamelCase							,							device=self.device							,							dtype=_UpperCamelCase	)
							else:
										if latents.shape != latents_shape:
													raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""	)
										lowerCAmelCase_         =      latents.to(self.device	)

							# set timesteps
							self.scheduler.set_timesteps(_UpperCamelCase	)

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

							# scale the initial noise by the standard deviation required by the scheduler
							lowerCAmelCase_         =      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]
							lowerCAmelCase_         =      "eta" in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
							lowerCAmelCase_         =      {}
							if accepts_eta:
										lowerCAmelCase_         =      eta

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

										# predict the noise residual
										lowerCAmelCase_         =      self.unet(_UpperCamelCase							,							_UpperCamelCase							,							encoder_hidden_states=_UpperCamelCase	).sample

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

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

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

							lowerCAmelCase_         =      1 / 0.18215 * latents
							lowerCAmelCase_         =      self.vae.decode(_UpperCamelCase	).sample

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

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

							if output_type == "pil":
										lowerCAmelCase_         =      self.numpy_to_pil(_UpperCamelCase	)

							if not return_dict:
										return image

							return StableDiffusionPipelineOutput(images=_UpperCamelCase							,							nsfw_content_detected=_UpperCamelCase	)





from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional

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


if TYPE_CHECKING:
			from ... import FeatureExtractionMixin, TensorType

_A  =	logging.get_logger(__name__)

_A  =	{
    "openai/imagegpt-small": "",
    "openai/imagegpt-medium": "",
    "openai/imagegpt-large": "",
}





class       _lowerCAmelCase   (      __a			):
				_lowercase     						='''imagegpt'''
				_lowercase     						=['''past_key_values''']
				_lowercase     						={
				    '''hidden_size''': '''n_embd''',
				    '''max_position_embeddings''': '''n_positions''',
				    '''num_attention_heads''': '''n_head''',
				    '''num_hidden_layers''': '''n_layer''',
				}




				def __init__(   self							,							_UpperCamelCase=512 + 1							,							_UpperCamelCase=32 * 32							,							_UpperCamelCase=512							,							_UpperCamelCase=24							,							_UpperCamelCase=8							,							_UpperCamelCase=None							,							_UpperCamelCase="quick_gelu"							,							_UpperCamelCase=0.1							,							_UpperCamelCase=0.1							,							_UpperCamelCase=0.1							,							_UpperCamelCase=1e-5							,							_UpperCamelCase=0.02							,							_UpperCamelCase=True							,							_UpperCamelCase=True							,							_UpperCamelCase=False							,							_UpperCamelCase=False							,							_UpperCamelCase=False							,							**_UpperCamelCase							,							)							-> Tuple:
							lowerCAmelCase_         =      vocab_size
							lowerCAmelCase_         =      n_positions
							lowerCAmelCase_         =      n_embd
							lowerCAmelCase_         =      n_layer
							lowerCAmelCase_         =      n_head
							lowerCAmelCase_         =      n_inner
							lowerCAmelCase_         =      activation_function
							lowerCAmelCase_         =      resid_pdrop
							lowerCAmelCase_         =      embd_pdrop
							lowerCAmelCase_         =      attn_pdrop
							lowerCAmelCase_         =      layer_norm_epsilon
							lowerCAmelCase_         =      initializer_range
							lowerCAmelCase_         =      scale_attn_weights
							lowerCAmelCase_         =      use_cache
							lowerCAmelCase_         =      scale_attn_by_inverse_layer_idx
							lowerCAmelCase_         =      reorder_and_upcast_attn
							lowerCAmelCase_         =      tie_word_embeddings

							super().__init__(tie_word_embeddings=_UpperCamelCase							,							**_UpperCamelCase	)





class       _lowerCAmelCase   (      __a			):




				@property
				def 	__a		(   self	)							-> Mapping[str, Mapping[int, str]]:
							return OrderedDict(
							    [
							        ("input_ids", {0: "batch", 1: "sequence"}),
							    ]	)




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase = 1							,							_UpperCamelCase = -1							,							_UpperCamelCase = False							,							_UpperCamelCase = None							,							_UpperCamelCase = 3							,							_UpperCamelCase = 32							,							_UpperCamelCase = 32							,							)							-> Mapping[str, Any]:
							lowerCAmelCase_         =      self._generate_dummy_images(_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)
							lowerCAmelCase_         =      dict(preprocessor(images=_UpperCamelCase							,							return_tensors=_UpperCamelCase	)	)

							return inputs



