





'''simple docstring'''






from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin







@dataclass
class 			lowerCamelCase					(						UpperCamelCase_  ):
					_SCREAMING_SNAKE_CASE							   =    42







class 			lowerCamelCase					(						UpperCamelCase_     ,							UpperCamelCase_  ):


					@register_to_config
					def __init__(					self	:  int       ,     __snake_case	:  int = 32       ,     __snake_case	:  int = 64       ,     __snake_case	:  int = 20       ,     __snake_case	:  int = 7_68       ,     __snake_case	:  Dict=77       ,     __snake_case	:  Any=4       ,     __snake_case	:  float = 0.0       ,     __snake_case	:  str = "silu"       ,     __snake_case	:  Optional[str] = None       ,     __snake_case	:  Optional[str] = None       ,     __snake_case	:  Optional[str] = "linear"       ,     __snake_case	:  Optional[str] = "prd"       ,     __snake_case	:  Optional[int] = None       ,     __snake_case	:  Optional[int] = None       ,     __snake_case	:  Optional[int] = None       ,     ):





											'''simple docstring'''





											super().__init__()
											_snake_case:   Dict  					=		num_attention_heads
											_snake_case:   Any  					=		attention_head_dim
											_snake_case:   Union[str, Any]  					=		num_attention_heads * attention_head_dim
											_snake_case:   Optional[Any]  					=		additional_embeddings

											_snake_case:   Tuple  					=		time_embed_dim or inner_dim
											_snake_case:   Optional[int]  					=		embedding_proj_dim or embedding_dim
											_snake_case:   Dict  					=		clip_embed_dim or embedding_dim

											_snake_case:   List[Any]  					=		Timesteps(__A       ,     __A       ,     0       )
											_snake_case:   int  					=		TimestepEmbedding(__A       ,     __A       ,     out_dim=__A       ,     act_fn=__A       )

											_snake_case:   List[str]  					=		nn.Linear(__A       ,     __A       )

											if embedding_proj_norm_type is None:
																	_snake_case:   int  					=		None
											elif embedding_proj_norm_type == "layer":
																	_snake_case:   int  					=		nn.LayerNorm(__A       )
											else:
																	raise ValueError(f'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'''       )

											_snake_case:   List[str]  					=		nn.Linear(__A       ,     __A       )

											if encoder_hid_proj_type is None:
																	_snake_case:   str  					=		None
											elif encoder_hid_proj_type == "linear":
																	_snake_case:   Optional[Any]  					=		nn.Linear(__A       ,     __A       )
											else:
																	raise ValueError(f'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'''       )

											_snake_case:   Any  					=		nn.Parameter(torch.zeros(1       ,     num_embeddings + additional_embeddings       ,     __A       )       )

											if added_emb_type == "prd":
																	_snake_case:   int  					=		nn.Parameter(torch.zeros(1       ,     1       ,     __A       )       )
											elif added_emb_type is None:
																	_snake_case:   Tuple  					=		None
											else:
																	raise ValueError(
																	    f'''`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'''       )

											_snake_case:   List[str]  					=		nn.ModuleList(
											    [
											        BasicTransformerBlock(
											            __A       ,     __A       ,     __A       ,     dropout=__A       ,     activation_fn='gelu'       ,     attention_bias=__A       ,     )
											        for d in range(__A       )
											    ]       )

											if norm_in_type == "layer":
																	_snake_case:   Any  					=		nn.LayerNorm(__A       )
											elif norm_in_type is None:
																	_snake_case:   str  					=		None
											else:
																	raise ValueError(f'''Unsupported norm_in_type: {norm_in_type}.'''       )

											_snake_case:   Optional[int]  					=		nn.LayerNorm(__A       )

											_snake_case:   Optional[Any]  					=		nn.Linear(__A       ,     __A       )

											_snake_case:   Optional[Any]  					=		torch.full(
											    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]       ,     -1_00_00.0       )
											causal_attention_mask.triu_(1       )
											_snake_case:   int  					=		causal_attention_mask[None, ...]
											self.register_buffer('causal_attention_mask'       ,     __A       ,     persistent=__A       )

											_snake_case:   Tuple  					=		nn.Parameter(torch.zeros(1       ,     __A       )       )
											_snake_case:   Union[str, Any]  					=		nn.Parameter(torch.zeros(1       ,     __A       )       )


					@property
					# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		{}

											def fn_recursive_add_processors(__snake_case	:  str       ,     __snake_case	:  torch.nn.Module       ,     __snake_case	:  Dict[str, AttentionProcessor]       ):
																	if hasattr(__A       ,     'set_processor'       ):
																							_snake_case:   List[str]  					=		module.processor

																	for sub_name, child in module.named_children():
																							fn_recursive_add_processors(f'''{name}.{sub_name}'''       ,     __A       ,     __A       )

																	return processors

											for name, module in self.named_children():
																	fn_recursive_add_processors(__A       ,     __A       ,     __A       )

											return processors


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ,     __snake_case	:  Union[AttentionProcessor, Dict[str, AttentionProcessor]]       ):





											'''simple docstring'''





											_snake_case:   int  					=		len(self.attn_processors.keys()       )

											if isinstance(__A       ,     __A       ) and len(__A       ) != count:
																	raise ValueError(
																	    f'''A dict of processors was passed, but the number of processors {len(__A       )} does not match the'''
																	    f''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''       )

											def fn_recursive_attn_processor(__snake_case	:  str       ,     __snake_case	:  torch.nn.Module       ,     __snake_case	:  Dict       ):
																	if hasattr(__A       ,     'set_processor'       ):
																							if not isinstance(__A       ,     __A       ):
																													module.set_processor(__A       )
																							else:
																													module.set_processor(processor.pop(f'''{name}.processor'''       )       )

																	for sub_name, child in module.named_children():
																							fn_recursive_attn_processor(f'''{name}.{sub_name}'''       ,     __A       ,     __A       )

											for name, module in self.named_children():
																	fn_recursive_attn_processor(__A       ,     __A       ,     __A       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											self.set_attn_processor(AttnProcessor()       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ,     __snake_case	:  Any       ,     __snake_case	:  Union[torch.Tensor, float, int]       ,     __snake_case	:  torch.FloatTensor       ,     __snake_case	:  Optional[torch.FloatTensor] = None       ,     __snake_case	:  Optional[torch.BoolTensor] = None       ,     __snake_case	:  bool = True       ,     ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		hidden_states.shape[0]

											_snake_case:   int  					=		timestep
											if not torch.is_tensor(__A       ):
																	_snake_case:   Dict  					=		torch.tensor([timesteps]       ,     dtype=torch.long       ,     device=hidden_states.device       )
											elif torch.is_tensor(__A       ) and len(timesteps.shape       ) == 0:
																	_snake_case:   Union[str, Any]  					=		timesteps[None].to(hidden_states.device       )

											# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
											_snake_case:   Optional[Any]  					=		timesteps * torch.ones(__A       ,     dtype=timesteps.dtype       ,     device=timesteps.device       )

											_snake_case:   Any  					=		self.time_proj(__A       )

											# timesteps does not contain any weights and will always return f32 tensors
											# but time_embedding might be fp16, so we need to cast here.
											_snake_case:   int  					=		timesteps_projected.to(dtype=self.dtype       )
											_snake_case:   List[str]  					=		self.time_embedding(__A       )

											if self.embedding_proj_norm is not None:
																	_snake_case:   List[str]  					=		self.embedding_proj_norm(__A       )

											_snake_case:   str  					=		self.embedding_proj(__A       )
											if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
																	_snake_case:   Any  					=		self.encoder_hidden_states_proj(__A       )
											elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
																	raise ValueError('`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'       )

											_snake_case:   List[str]  					=		self.proj_in(__A       )

											_snake_case:   Union[str, Any]  					=		self.positional_embedding.to(hidden_states.dtype       )

											_snake_case:   Optional[int]  					=		[]
											_snake_case:   Dict  					=		0

											if encoder_hidden_states is not None:
																	additional_embeds.append(__A       )
																	additional_embeddings_len += encoder_hidden_states.shape[1]

											if len(proj_embeddings.shape       ) == 2:
																	_snake_case:   Tuple  					=		proj_embeddings[:, None, :]

											if len(hidden_states.shape       ) == 2:
																	_snake_case:   Dict  					=		hidden_states[:, None, :]

											_snake_case:   Any  					=		additional_embeds + [
											    proj_embeddings,
											    time_embeddings[:, None, :],
											    hidden_states,
											]

											if self.prd_embedding is not None:
																	_snake_case:   Any  					=		self.prd_embedding.to(hidden_states.dtype       ).expand(__A       ,     -1       ,     -1       )
																	additional_embeds.append(__A       )

											_snake_case:   Dict  					=		torch.cat(
											    __A       ,     dim=1       ,     )

											# Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
											_snake_case:   List[Any]  					=		additional_embeddings_len + proj_embeddings.shape[1] + 1
											if positional_embeddings.shape[1] < hidden_states.shape[1]:
																	_snake_case:   Union[str, Any]  					=		F.pad(
																	    __A       ,     (
																	        0,
																	        0,
																	        additional_embeddings_len,
																	        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
																	    )       ,     value=0.0       ,     )

											_snake_case:   List[Any]  					=		hidden_states + positional_embeddings

											if attention_mask is not None:
																	_snake_case:   Any  					=		(1 - attention_mask.to(hidden_states.dtype       )) * -1_00_00.0
																	_snake_case:   Tuple  					=		F.pad(__A       ,     (0, self.additional_embeddings)       ,     value=0.0       )
																	_snake_case:   Dict  					=		(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype       )
																	_snake_case:   List[str]  					=		attention_mask.repeat_interleave(self.config.num_attention_heads       ,     dim=0       )

											if self.norm_in is not None:
																	_snake_case:   Optional[int]  					=		self.norm_in(__A       )

											for block in self.transformer_blocks:
																	_snake_case:   Union[str, Any]  					=		block(__A       ,     attention_mask=__A       )

											_snake_case:   Optional[int]  					=		self.norm_out(__A       )

											if self.prd_embedding is not None:
																	_snake_case:   Dict  					=		hidden_states[:, -1]
											else:
																	_snake_case:   Optional[int]  					=		hidden_states[:, additional_embeddings_len:]

											_snake_case:   Union[str, Any]  					=		self.proj_to_clip_embeddings(__A       )

											if not return_dict:
																	return (predicted_image_embedding,)

											return PriorTransformerOutput(predicted_image_embedding=__A       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ,     __snake_case	:  Optional[int]       ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		(prior_latents * self.clip_std) + self.clip_mean
											return prior_latents












'''simple docstring'''






import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin







class 			lowerCamelCase					(						__UpperCAmelCase  ):
					_SCREAMING_SNAKE_CASE							   =    "Speech2TextFeatureExtractor"
					_SCREAMING_SNAKE_CASE							   =    "Speech2TextTokenizer"


					def __init__(					self	:  List[str]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  List[str]       ):





											'''simple docstring'''





											super().__init__(__snake_case       ,     __snake_case       )
											_snake_case:   Optional[int]  					=		self.feature_extractor
											_snake_case:   Dict  					=		False


					def __call__(					self	:  Optional[int]       ,     *__snake_case	:  Optional[Any]       ,     **__snake_case	:  Union[str, Any]       ):





											'''simple docstring'''





											if self._in_target_context_manager:
																	return self.current_processor(*__snake_case       ,     **__snake_case       )

											if "raw_speech" in kwargs:
																	warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'       )
																	_snake_case:   int  					=		kwargs.pop('raw_speech'       )
											else:
																	_snake_case:   Tuple  					=		kwargs.pop('audio'       ,     __snake_case       )
											_snake_case:   Optional[Any]  					=		kwargs.pop('sampling_rate'       ,     __snake_case       )
											_snake_case:   Optional[Any]  					=		kwargs.pop('text'       ,     __snake_case       )
											if len(__snake_case       ) > 0:
																	_snake_case:   Optional[Any]  					=		args[0]
																	_snake_case:   Union[str, Any]  					=		args[1:]

											if audio is None and text is None:
																	raise ValueError('You need to specify either an `audio` or `text` input to process.'       )

											if audio is not None:
																	_snake_case:   List[Any]  					=		self.feature_extractor(__snake_case       ,     *__snake_case       ,     sampling_rate=__snake_case       ,     **__snake_case       )
											if text is not None:
																	_snake_case:   Tuple  					=		self.tokenizer(__snake_case       ,     **__snake_case       )

											if text is None:
																	return inputs
											elif audio is None:
																	return encodings
											else:
																	_snake_case:   List[Any]  					=		encodings['input_ids']
																	return inputs


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     *__snake_case	:  Optional[int]       ,     **__snake_case	:  str       ):





											'''simple docstring'''





											return self.tokenizer.batch_decode(*__snake_case       ,     **__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     *__snake_case	:  List[str]       ,     **__snake_case	:  Tuple       ):





											'''simple docstring'''





											return self.tokenizer.decode(*__snake_case       ,     **__snake_case       )




					@contextmanager
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											warnings.warn(
											    '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
											    'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
											    'your audio inputs, or in a separate call.'       )
											_snake_case:   Any  					=		True
											_snake_case:   str  					=		self.tokenizer
											yield
											_snake_case:   Union[str, Any]  					=		self.feature_extractor
											_snake_case:   Tuple  					=		False





