



from ...processing_utils import ProcessorMixin






class    SCREAMING_SNAKE_CASE			(							_UpperCamelCase					):
       __lowerCamelCase	:	Tuple    ='WhisperFeatureExtractor'
       __lowerCamelCase	:	Optional[Any]    ='WhisperTokenizer'





       def __init__(   self		:		List[Any]      ,				__lowercase		:		int      ,				__lowercase		:		Optional[int]	):

           '''simple docstring'''
           super().__init__(_UpperCAmelCase      ,				_UpperCAmelCase	)
           __a       =		self.feature_extractor
           __a       =		False





       def 						UpperCamelCase_					(   self		:		Any      ,				__lowercase		:		Tuple=None      ,				__lowercase		:		Tuple=None      ,				__lowercase		:		Dict=True	):

           '''simple docstring'''
           return self.tokenizer.get_decoder_prompt_ids(task=_UpperCAmelCase      ,				language=_UpperCAmelCase      ,				no_timestamps=_UpperCAmelCase	)





       def __call__(   self		:		Optional[Any]      ,				*__lowercase		:		str      ,				**__lowercase		:		str	):

           '''simple docstring'''
           # For backward compatibility
           if self._in_target_context_manager:
               return self.current_processor(*_UpperCAmelCase      ,				**_UpperCAmelCase	)

           __a       =		kwargs.pop("""audio"""      ,				_UpperCAmelCase	)
           __a       =		kwargs.pop("""sampling_rate"""      ,				_UpperCAmelCase	)
           __a       =		kwargs.pop("""text"""      ,				_UpperCAmelCase	)
           if len(_UpperCAmelCase	) > 0:
               __a       =		args[0]
               __a       =		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:
               __a       =		self.feature_extractor(_UpperCAmelCase      ,				*_UpperCAmelCase      ,				sampling_rate=_UpperCAmelCase      ,				**_UpperCAmelCase	)
           if text is not None:
               __a       =		self.tokenizer(_UpperCAmelCase      ,				**_UpperCAmelCase	)

           if text is None:
               return inputs

           elif audio is None:
               return encodings
           else:
               __a       =		encodings['''input_ids''']
               return inputs





       def 						UpperCamelCase_					(   self		:		Any      ,				*__lowercase		:		str      ,				**__lowercase		:		Optional[Any]	):

           '''simple docstring'''
           return self.tokenizer.batch_decode(*_UpperCAmelCase      ,				**_UpperCAmelCase	)





       def 						UpperCamelCase_					(   self		:		Dict      ,				*__lowercase		:		Union[str, Any]      ,				**__lowercase		:		Any	):

           '''simple docstring'''
           return self.tokenizer.decode(*_UpperCAmelCase      ,				**_UpperCAmelCase	)







       def 						UpperCamelCase_					(   self		:		Union[str, Any]      ,				__lowercase		:		Optional[int]      ,				__lowercase		:		Optional[Any]="np"	):

           '''simple docstring'''
           return self.tokenizer.get_prompt_ids(_UpperCAmelCase      ,				return_tensors=_UpperCAmelCase	)





from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__					):
							__lowerCamelCase	:	torch.FloatTensor






class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__			,      lowerCamelCase__					):





							@register_to_config
							def __init__(   self		:		Tuple      ,				__lowercase		:		int = 16      ,				__lowercase		:		int = 88      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		int = 1      ,				__lowercase		:		float = 0.0      ,				__lowercase		:		int = 32      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		bool = False      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		str = "geglu"      ,				__lowercase		:		bool = True      ,				__lowercase		:		bool = True      ,				):

											'''simple docstring'''
											super().__init__()
											__a       =		num_attention_heads
											__a       =		attention_head_dim
											__a       =		num_attention_heads * attention_head_dim

											__a       =		in_channels

											__a       =		torch.nn.GroupNorm(num_groups=__lowercase      ,				num_channels=__lowercase      ,				eps=1E-6      ,				affine=__lowercase	)
											__a       =		nn.Linear(__lowercase      ,				__lowercase	)

											# 3. Define transformers blocks
											__a       =		nn.ModuleList(
											    [
											        BasicTransformerBlock(
											            __lowercase      ,				__lowercase      ,				__lowercase      ,				dropout=__lowercase      ,				cross_attention_dim=__lowercase      ,				activation_fn=__lowercase      ,				attention_bias=__lowercase      ,				double_self_attention=__lowercase      ,				norm_elementwise_affine=__lowercase      ,				)
											        for d in range(__lowercase	)
											    ]	)

											__a       =		nn.Linear(__lowercase      ,				__lowercase	)







							def 						UpperCamelCase_					(   self		:		Any      ,				__lowercase		:		Optional[Any]      ,				__lowercase		:		str=None      ,				__lowercase		:		List[Any]=None      ,				__lowercase		:		List[Any]=None      ,				__lowercase		:		int=1      ,				__lowercase		:		Union[str, Any]=None      ,				__lowercase		:		bool = True      ,				):

											'''simple docstring'''
											__a     ,					__a     ,					__a     ,					__a       =		hidden_states.shape
											__a       =		batch_frames // num_frames

											__a       =		hidden_states

											__a       =		hidden_states[None, :].reshape(__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase	)
											__a       =		hidden_states.permute(0      ,				2      ,				1      ,				3      ,				4	)

											__a       =		self.norm(__lowercase	)
											__a       =		hidden_states.permute(0      ,				3      ,				4      ,				2      ,				1	).reshape(batch_size * height * width      ,				__lowercase      ,				__lowercase	)

											__a       =		self.proj_in(__lowercase	)

											# 2. Blocks
											for block in self.transformer_blocks:
															__a       =		block(
															    __lowercase      ,				encoder_hidden_states=__lowercase      ,				timestep=__lowercase      ,				cross_attention_kwargs=__lowercase      ,				class_labels=__lowercase      ,				)

											# 3. Output
											__a       =		self.proj_out(__lowercase	)
											__a       =		(
											    hidden_states[None, None, :]
											    .reshape(__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase	)
											    .permute(0      ,				3      ,				4      ,				1      ,				2	)
											    .contiguous()
											)
											__a       =		hidden_states.reshape(__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase	)

											__a       =		hidden_states + residual

											if not return_dict:
															return (output,)

											return TransformerTemporalModelOutput(sample=__lowercase	)
