"""simple docstring"""




from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block
@dataclass
class     UpperCamelCase__(			__A ):
						lowerCAmelCase__  :						torch.FloatTensor






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase=3		,__UpperCAmelCase=3		,__UpperCAmelCase=("DownEncoderBlock2D",)		,__UpperCAmelCase=(64,)		,__UpperCAmelCase=2		,__UpperCAmelCase=32		,__UpperCAmelCase="silu"		,__UpperCAmelCase=True		,)					->				Optional[int]:
													super().__init__()
													A__	     =							layers_per_block

													A__	     =							torch.nn.Convad(
													    __UpperCAmelCase		,block_out_channels[0]		,kernel_size=3		,stride=1		,padding=1		,)

													A__	     =							None
													A__	     =							nn.ModuleList([]     )

													# down
													A__	     =							block_out_channels[0]
													for i, down_block_type in enumerate(__UpperCAmelCase     ):
																				A__	     =							output_channel
																				A__	     =							block_out_channels[i]
																				A__	     =							i == len(__UpperCAmelCase     ) - 1

																				A__	     =							get_down_block(
																				    __UpperCAmelCase		,num_layers=self.layers_per_block		,in_channels=__UpperCAmelCase		,out_channels=__UpperCAmelCase		,add_downsample=not is_final_block		,resnet_eps=1e-6		,downsample_padding=0		,resnet_act_fn=__UpperCAmelCase		,resnet_groups=__UpperCAmelCase		,attention_head_dim=__UpperCAmelCase		,temb_channels=__UpperCAmelCase		,)
																				self.down_blocks.append(__UpperCAmelCase     )

													# mid
													A__	     =							UNetMidBlockaD(
													    in_channels=block_out_channels[-1]		,resnet_eps=1e-6		,resnet_act_fn=__UpperCAmelCase		,output_scale_factor=1		,resnet_time_scale_shift='default'		,attention_head_dim=block_out_channels[-1]		,resnet_groups=__UpperCAmelCase		,temb_channels=__UpperCAmelCase		,)

													# out
													A__	     =							nn.GroupNorm(num_channels=block_out_channels[-1]		,num_groups=__UpperCAmelCase		,eps=1e-6     )
													A__	     =							nn.SiLU()

													A__	     =							2 * out_channels if double_z else out_channels
													A__	     =							nn.Convad(block_out_channels[-1]		,__UpperCAmelCase		,3		,padding=1     )

													A__	     =							False





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				List[str]:
													A__	     =							x
													A__	     =							self.conv_in(__UpperCAmelCase     )

													if self.training and self.gradient_checkpointing:

																				def create_custom_forward(__UpperCAmelCase     ):
																											def custom_forward(*__UpperCAmelCase     ):
																																		return module(*__UpperCAmelCase     )

																											return custom_forward

																				# down
																				if is_torch_version('>='		,'1.11.0'     ):
																											for down_block in self.down_blocks:
																																		A__	     =							torch.utils.checkpoint.checkpoint(
																																		    create_custom_forward(__UpperCAmelCase     )		,__UpperCAmelCase		,use_reentrant=__UpperCAmelCase     )
																											# middle
																											A__	     =							torch.utils.checkpoint.checkpoint(
																											    create_custom_forward(self.mid_block     )		,__UpperCAmelCase		,use_reentrant=__UpperCAmelCase     )
																				else:
																											for down_block in self.down_blocks:
																																		A__	     =							torch.utils.checkpoint.checkpoint(create_custom_forward(__UpperCAmelCase     )		,__UpperCAmelCase     )
																											# middle
																											A__	     =							torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block     )		,__UpperCAmelCase     )

													else:
																				# down
																				for down_block in self.down_blocks:
																											A__	     =							down_block(__UpperCAmelCase     )

																				# middle
																				A__	     =							self.mid_block(__UpperCAmelCase     )

													# post-process
													A__	     =							self.conv_norm_out(__UpperCAmelCase     )
													A__	     =							self.conv_act(__UpperCAmelCase     )
													A__	     =							self.conv_out(__UpperCAmelCase     )

													return sample






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase=3		,__UpperCAmelCase=3		,__UpperCAmelCase=("UpDecoderBlock2D",)		,__UpperCAmelCase=(64,)		,__UpperCAmelCase=2		,__UpperCAmelCase=32		,__UpperCAmelCase="silu"		,__UpperCAmelCase="group"		,)					->				Optional[int]:
													super().__init__()
													A__	     =							layers_per_block

													A__	     =							nn.Convad(
													    __UpperCAmelCase		,block_out_channels[-1]		,kernel_size=3		,stride=1		,padding=1		,)

													A__	     =							None
													A__	     =							nn.ModuleList([]     )

													A__	     =							in_channels if norm_type == 'spatial' else None

													# mid
													A__	     =							UNetMidBlockaD(
													    in_channels=block_out_channels[-1]		,resnet_eps=1e-6		,resnet_act_fn=__UpperCAmelCase		,output_scale_factor=1		,resnet_time_scale_shift='default' if norm_type == 'group' else norm_type		,attention_head_dim=block_out_channels[-1]		,resnet_groups=__UpperCAmelCase		,temb_channels=__UpperCAmelCase		,)

													# up
													A__	     =							list(reversed(__UpperCAmelCase     )     )
													A__	     =							reversed_block_out_channels[0]
													for i, up_block_type in enumerate(__UpperCAmelCase     ):
																				A__	     =							output_channel
																				A__	     =							reversed_block_out_channels[i]

																				A__	     =							i == len(__UpperCAmelCase     ) - 1

																				A__	     =							get_up_block(
																				    __UpperCAmelCase		,num_layers=self.layers_per_block + 1		,in_channels=__UpperCAmelCase		,out_channels=__UpperCAmelCase		,prev_output_channel=__UpperCAmelCase		,add_upsample=not is_final_block		,resnet_eps=1e-6		,resnet_act_fn=__UpperCAmelCase		,resnet_groups=__UpperCAmelCase		,attention_head_dim=__UpperCAmelCase		,temb_channels=__UpperCAmelCase		,resnet_time_scale_shift=__UpperCAmelCase		,)
																				self.up_blocks.append(__UpperCAmelCase     )
																				A__	     =							output_channel

													# out
													if norm_type == "spatial":
																				A__	     =							SpatialNorm(block_out_channels[0]		,__UpperCAmelCase     )
													else:
																				A__	     =							nn.GroupNorm(num_channels=block_out_channels[0]		,num_groups=__UpperCAmelCase		,eps=1e-6     )
													A__	     =							nn.SiLU()
													A__	     =							nn.Convad(block_out_channels[0]		,__UpperCAmelCase		,3		,padding=1     )

													A__	     =							False





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase=None     )					->				Dict:
													A__	     =							z
													A__	     =							self.conv_in(__UpperCAmelCase     )

													A__	     =							next(iter(self.up_blocks.parameters()     )     ).dtype
													if self.training and self.gradient_checkpointing:

																				def create_custom_forward(__UpperCAmelCase     ):
																											def custom_forward(*__UpperCAmelCase     ):
																																		return module(*__UpperCAmelCase     )

																											return custom_forward

																				if is_torch_version('>='		,'1.11.0'     ):
																											# middle
																											A__	     =							torch.utils.checkpoint.checkpoint(
																											    create_custom_forward(self.mid_block     )		,__UpperCAmelCase		,__UpperCAmelCase		,use_reentrant=__UpperCAmelCase     )
																											A__	     =							sample.to(__UpperCAmelCase     )

																											# up
																											for up_block in self.up_blocks:
																																		A__	     =							torch.utils.checkpoint.checkpoint(
																																		    create_custom_forward(__UpperCAmelCase     )		,__UpperCAmelCase		,__UpperCAmelCase		,use_reentrant=__UpperCAmelCase     )
																				else:
																											# middle
																											A__	     =							torch.utils.checkpoint.checkpoint(
																											    create_custom_forward(self.mid_block     )		,__UpperCAmelCase		,__UpperCAmelCase     )
																											A__	     =							sample.to(__UpperCAmelCase     )

																											# up
																											for up_block in self.up_blocks:
																																		A__	     =							torch.utils.checkpoint.checkpoint(create_custom_forward(__UpperCAmelCase     )		,__UpperCAmelCase		,__UpperCAmelCase     )
													else:
																				# middle
																				A__	     =							self.mid_block(__UpperCAmelCase		,__UpperCAmelCase     )
																				A__	     =							sample.to(__UpperCAmelCase     )

																				# up
																				for up_block in self.up_blocks:
																											A__	     =							up_block(__UpperCAmelCase		,__UpperCAmelCase     )

        # post-process
													if latent_embeds is None:
																				A__	     =							self.conv_norm_out(__UpperCAmelCase     )
													else:
																				A__	     =							self.conv_norm_out(__UpperCAmelCase		,__UpperCAmelCase     )
													A__	     =							self.conv_act(__UpperCAmelCase     )
													A__	     =							self.conv_out(__UpperCAmelCase     )

													return sample






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=None		,__UpperCAmelCase="random"		,__UpperCAmelCase=False		,__UpperCAmelCase=True     )					->				Optional[Any]:
													super().__init__()
													A__	     =							n_e
													A__	     =							vq_embed_dim
													A__	     =							beta
													A__	     =							legacy

													A__	     =							nn.Embedding(self.n_e		,self.vq_embed_dim     )
													self.embedding.weight.data.uniform_(-1.0 / self.n_e		,1.0 / self.n_e     )

													A__	     =							remap
													if self.remap is not None:
																				self.register_buffer('used'		,torch.tensor(np.load(self.remap     )     )     )
																				A__	     =							self.used.shape[0]
																				A__	     =							unknown_index  # "random" or "extra" or integer
																				if self.unknown_index == "extra":
																											A__	     =							self.re_embed
																											A__	     =							self.re_embed + 1
																				print(
																				    f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
																				    f'''Using {self.unknown_index} for unknown indices.'''     )
													else:
																				A__	     =							n_e

													A__	     =							sane_index_shape





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Tuple:
													A__	     =							inds.shape
													assert len(__UpperCAmelCase     ) > 1
													A__	     =							inds.reshape(ishape[0]		,-1     )
													A__	     =							self.used.to(__UpperCAmelCase     )
													A__	     =							(inds[:, :, None] == used[None, None, ...]).long()
													A__	     =							match.argmax(-1     )
													A__	     =							match.sum(2     ) < 1
													if self.unknown_index == "random":
																				A__	     =							torch.randint(0		,self.re_embed		,size=new[unknown].shape     ).to(device=new.device     )
													else:
																				A__	     =							self.unknown_index
													return new.reshape(__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				str:
													A__	     =							inds.shape
													assert len(__UpperCAmelCase     ) > 1
													A__	     =							inds.reshape(ishape[0]		,-1     )
													A__	     =							self.used.to(__UpperCAmelCase     )
													if self.re_embed > self.used.shape[0]:  # extra token
																				A__	     =							0  # simply set to zero
													A__	     =							torch.gather(used[None, :][inds.shape[0] * [0], :]		,1		,__UpperCAmelCase     )
													return back.reshape(__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				int:
													# reshape z -> (batch, height, width, channel) and flatten
													A__	     =							z.permute(0		,2		,3		,1     ).contiguous()
													A__	     =							z.view(-1		,self.vq_embed_dim     )

													# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
													A__	     =							torch.argmin(torch.cdist(__UpperCAmelCase		,self.embedding.weight     )		,dim=1     )

													A__	     =							self.embedding(__UpperCAmelCase     ).view(z.shape     )
													A__	     =							None
													A__	     =							None

													# compute loss for embedding
													if not self.legacy:
																				A__	     =							self.beta * torch.mean((z_q.detach() - z) ** 2     ) + torch.mean((z_q - z.detach()) ** 2     )
													else:
																				A__	     =							torch.mean((z_q.detach() - z) ** 2     ) + self.beta * torch.mean((z_q - z.detach()) ** 2     )

													# preserve gradients
													A__	     =							z + (z_q - z).detach()

													# reshape back to match original input shape
													A__	     =							z_q.permute(0		,3		,1		,2     ).contiguous()

													if self.remap is not None:
																				A__	     =							min_encoding_indices.reshape(z.shape[0]		,-1     )  # add batch axis
																				A__	     =							self.remap_to_used(__UpperCAmelCase     )
																				A__	     =							min_encoding_indices.reshape(-1		,1     )  # flatten

													if self.sane_index_shape:
																				A__	     =							min_encoding_indices.reshape(z_q.shape[0]		,z_q.shape[2]		,z_q.shape[3]     )

													return z_q, loss, (perplexity, min_encodings, min_encoding_indices)





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase     )					->				Any:
													# shape specifying (batch, height, width, channel)
													if self.remap is not None:
																				A__	     =							indices.reshape(shape[0]		,-1     )  # add batch axis
																				A__	     =							self.unmap_to_all(__UpperCAmelCase     )
																				A__	     =							indices.reshape(-1     )  # flatten again

													# get quantized latent vectors
													A__	     =							self.embedding(__UpperCAmelCase     )

													if shape is not None:
																				A__	     =							z_q.view(__UpperCAmelCase     )
																				# reshape back to match original input shape
																				A__	     =							z_q.permute(0		,3		,1		,2     ).contiguous()

													return z_q






class     UpperCamelCase__(			__A ):





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase=False     )					->				Union[str, Any]:
													A__	     =							parameters
													A__     ,  A__	     =							torch.chunk(__UpperCAmelCase		,2		,dim=1     )
													A__	     =							torch.clamp(self.logvar		,-3_0.0		,2_0.0     )
													A__	     =							deterministic
													A__	     =							torch.exp(0.5 * self.logvar     )
													A__	     =							torch.exp(self.logvar     )
													if self.deterministic:
																				A__	     =							A__	     =							torch.zeros_like(
																				    self.mean		,device=self.parameters.device		,dtype=self.parameters.dtype     )





						def 				snake_case__	(       self		,__UpperCAmelCase = None     )					->				torch.FloatTensor:
													# make sure sample is on the same device as the parameters and has same dtype
													A__	     =							randn_tensor(
													    self.mean.shape		,generator=__UpperCAmelCase		,device=self.parameters.device		,dtype=self.parameters.dtype     )
													A__	     =							self.mean + self.std * sample
													return x





						def 				snake_case__	(       self		,__UpperCAmelCase=None     )					->				int:
													if self.deterministic:
																				return torch.Tensor([0.0]     )
													else:
																				if other is None:
																											return 0.5 * torch.sum(torch.pow(self.mean		,2     ) + self.var - 1.0 - self.logvar		,dim=[1, 2, 3]     )
																				else:
																											return 0.5 * torch.sum(
																											    torch.pow(self.mean - other.mean		,2     ) / other.var
																											    + self.var / other.var
																											    - 1.0
																											    - self.logvar
																											    + other.logvar		,dim=[1, 2, 3]		,)





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase=[1, 2, 3]     )					->				Union[str, Any]:
													if self.deterministic:
																				return torch.Tensor([0.0]     )
													A__	     =							np.log(2.0 * np.pi     )
													return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean		,2     ) / self.var		,dim=__UpperCAmelCase     )





						def 				snake_case__	(       self     )					->				Any:
													return self.mean






"""simple docstring"""




from io import BytesIO
from typing import List, Union

import requests

from ..utils import add_end_docstrings, is_decord_available, is_torch_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_decord_available():
				import numpy as np
				from decord import VideoReader


if is_torch_available():
				from ..models.auto.modeling_auto import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING

__lowerCamelCase											=  logging.get_logger(__name__)
@add_end_docstrings(__A )
class     UpperCamelCase__(			__A ):





						def __init__(       self		,*__UpperCAmelCase		,**__UpperCAmelCase     )					->				List[Any]:
													super().__init__(*__UpperCAmelCase		,**__UpperCAmelCase     )
													requires_backends(self		,'decord'     )
													self.check_model_type(__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase=None		,__UpperCAmelCase=None		,__UpperCAmelCase=None     )					->				Any:
													A__	     =							{}
													if frame_sampling_rate is not None:
																				A__	     =							frame_sampling_rate
													if num_frames is not None:
																				A__	     =							num_frames

													A__	     =							{}
													if top_k is not None:
																				A__	     =							top_k
													return preprocess_params, {}, postprocess_params





						def __call__(       self		,__UpperCAmelCase		,**__UpperCAmelCase     )					->				Optional[int]:
													return super().__call__(__UpperCAmelCase		,**__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase=None		,__UpperCAmelCase=1     )					->				List[Any]:
													if num_frames is None:
																				A__	     =							self.model.config.num_frames

													if video.startswith('http://'     ) or video.startswith('https://'     ):
																				A__	     =							BytesIO(requests.get(__UpperCAmelCase     ).content     )

													A__	     =							VideoReader(__UpperCAmelCase     )
													videoreader.seek(0     )

													A__	     =							0
													A__	     =							num_frames * frame_sampling_rate - 1
													A__	     =							np.linspace(__UpperCAmelCase		,__UpperCAmelCase		,num=__UpperCAmelCase		,dtype=np.intaa     )

													A__	     =							videoreader.get_batch(__UpperCAmelCase     ).asnumpy()
													A__	     =							list(__UpperCAmelCase     )

													A__	     =							self.image_processor(__UpperCAmelCase		,return_tensors=self.framework     )
													return model_inputs





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Optional[int]:
													A__	     =							self.model(**__UpperCAmelCase     )
													return model_outputs





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase=5     )					->				int:
													if top_k > self.model.config.num_labels:
																				A__	     =							self.model.config.num_labels

													if self.framework == "pt":
																				A__	     =							model_outputs.logits.softmax(-1     )[0]
																				A__     ,  A__	     =							probs.topk(__UpperCAmelCase     )
													else:
																				raise ValueError(f'''Unsupported framework: {self.framework}'''     )

													A__	     =							scores.tolist()
													A__	     =							ids.tolist()
													return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(__UpperCAmelCase		,__UpperCAmelCase     )]





