


"""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 			lowerCamelCase__						(   __snake_case			):
						lowerCAmelCase       =     42



class 			lowerCamelCase__						(   nn.Module			):







						def __init__(							self   :    int			,			_lowercase   :    int=3			,			_lowercase   :    Tuple=3			,			_lowercase   :    str=("DownEncoderBlock2D",)			,			_lowercase   :    Dict=(64,)			,			_lowercase   :    Optional[Any]=2			,			_lowercase   :    Union[str, Any]=32			,			_lowercase   :    List[Any]="silu"			,			_lowercase   :    Union[str, Any]=True			,			):
								super().__init__()
								A					       = layers_per_block

								A					       = torch.nn.Convad(
								    A_			,			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(A_		):
										A					       = output_channel
										A					       = block_out_channels[i]
										A					       = i == len(A_		) - 1

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

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

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

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

								A					       = False







						def __a						(							self   :    Tuple			,			_lowercase   :    int		):
								A					       = x
								A					       = self.conv_in(A_		)

								if self.training and self.gradient_checkpointing:

										def create_custom_forward(_lowercase   :    Optional[int]		):
												def custom_forward(*_lowercase   :    List[str]		):
														return module(*A_		)

												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(A_		)			,			A_			,			use_reentrant=A_		)
												# middle
												A					       = torch.utils.checkpoint.checkpoint(
												    create_custom_forward(self.mid_block		)			,			A_			,			use_reentrant=A_		)
										else:
												for down_block in self.down_blocks:
														A					       = torch.utils.checkpoint.checkpoint(create_custom_forward(A_		)			,			A_		)
												# middle
												A					       = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block		)			,			A_		)

								else:
										# down
										for down_block in self.down_blocks:
												A					       = down_block(A_		)

										# middle
										A					       = self.mid_block(A_		)

								# post-process
								A					       = self.conv_norm_out(A_		)
								A					       = self.conv_act(A_		)
								A					       = self.conv_out(A_		)

								return sample



class 			lowerCamelCase__						(   nn.Module			):







						def __init__(							self   :    Optional[Any]			,			_lowercase   :    Any=3			,			_lowercase   :    Union[str, Any]=3			,			_lowercase   :    Optional[Any]=("UpDecoderBlock2D",)			,			_lowercase   :    Tuple=(64,)			,			_lowercase   :    str=2			,			_lowercase   :    Dict=32			,			_lowercase   :    Any="silu"			,			_lowercase   :    Any="group"			,			):
								super().__init__()
								A					       = layers_per_block

								A					       = nn.Convad(
								    A_			,			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=A_			,			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=A_			,			temb_channels=A_			,			)

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

										A					       = i == len(A_		) - 1

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

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

								A					       = False







						def __a						(							self   :    str			,			_lowercase   :    int			,			_lowercase   :    List[Any]=None		):
								A					       = z
								A					       = self.conv_in(A_		)

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

										def create_custom_forward(_lowercase   :    str		):
												def custom_forward(*_lowercase   :    int		):
														return module(*A_		)

												return custom_forward

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

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

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

										# up
										for up_block in self.up_blocks:
												A					       = up_block(A_			,			A_		)

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

								return sample



class 			lowerCamelCase__						(   nn.Module			):







						def __init__(							self   :    Optional[Any]			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Dict			,			_lowercase   :    Optional[int]			,			_lowercase   :    List[Any]=None			,			_lowercase   :    List[Any]="random"			,			_lowercase   :    List[str]=False			,			_lowercase   :    str=True		):
								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 __a						(							self   :    Optional[int]			,			_lowercase   :    Dict		):
								A					       = inds.shape
								assert len(A_		) > 1
								A					       = inds.reshape(ishape[0]			,			-1		)
								A					       = self.used.to(A_		)
								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(A_		)







						def __a						(							self   :    Any			,			_lowercase   :    Any		):
								A					       = inds.shape
								assert len(A_		) > 1
								A					       = inds.reshape(ishape[0]			,			-1		)
								A					       = self.used.to(A_		)
								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			,			A_		)
								return back.reshape(A_		)







						def __a						(							self   :    Optional[Any]			,			_lowercase   :    Optional[int]		):
								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(A_			,			self.embedding.weight		)			,			dim=1		)

								A					       = self.embedding(A_		).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(A_		)
										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 __a						(							self   :    Any			,			_lowercase   :    List[str]			,			_lowercase   :    Optional[int]		):




								if self.remap is not None:
										A					       = indices.reshape(shape[0]			,			-1		)  # add batch axis
										A					       = self.unmap_to_all(A_		)
										A					       = indices.reshape(-1		)  # flatten again

								# get quantized latent vectors
								A					       = self.embedding(A_		)

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

								return z_q



class 			lowerCamelCase__						(   __snake_case			):







						def __init__(							self   :    str			,			_lowercase   :    int			,			_lowercase   :    str=False		):
								A					       = parameters
								A					       = torch.chunk(A_			,			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					       = torch.zeros_like(
										    self.mean			,			device=self.parameters.device			,			dtype=self.parameters.dtype		)







						def __a						(							self   :    Tuple			,			_lowercase   :    Dict = None		):
								A					       = randn_tensor(
								    self.mean.shape			,			generator=A_			,			device=self.parameters.device			,			dtype=self.parameters.dtype		)
								A					       = self.mean + self.std * sample
								return x







						def __a						(							self   :    Union[str, Any]			,			_lowercase   :    Union[str, Any]=None		):




								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 __a						(							self   :    Union[str, Any]			,			_lowercase   :    str			,			_lowercase   :    Union[str, Any]=[1, 2, 3]		):




								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=A_		)







						def __a						(							self   :    Optional[int]		):
								return self.mean









"""simple docstring"""


import unittest

from huggingface_hub import hf_hub_download

from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VideoMAEFeatureExtractor
from transformers.pipelines import VideoClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_decord,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
)

from .test_pipelines_common import ANY



@is_pipeline_test
@require_torch_or_tf
@require_vision
@require_decord
class 			lowerCamelCase__						(   unittest.TestCase			):
						lowerCAmelCase       =     MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING







						def __a						(							self   :    Dict			,			_lowercase   :    int			,			_lowercase   :    Any			,			_lowercase   :    int		):
								A					       = hf_hub_download(
								    repo_id='nateraw/video-demo'			,			filename='archery.mp4'			,			repo_type='dataset'		)
								A					       = VideoClassificationPipeline(model=_lowercase			,			image_processor=_lowercase			,			top_k=2		)
								A					       = [
								    example_video_filepath,
								    'https://huggingface.co/datasets/nateraw/video-demo/resolve/main/archery.mp4',
								]
								return video_classifier, examples







						def __a						(							self   :    Tuple			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    List[Any]		):
								for example in examples:
										A					       = video_classifier(_lowercase		)

										self.assertEqual(
										    _lowercase			,			[
										        {'score': ANY(_lowercase		), 'label': ANY(_lowercase		)},
										        {'score': ANY(_lowercase		), 'label': ANY(_lowercase		)},
										    ]			,			)







						@require_torch
						def __a						(							self   :    str		):
								A					       = 'hf-internal-testing/tiny-random-VideoMAEForVideoClassification'
								A					       = VideoMAEFeatureExtractor(
								    size={'shortest_edge': 10}			,			crop_size={'height': 10, 'width': 10}		)
								A					       = pipeline(
								    'video-classification'			,			model=_lowercase			,			feature_extractor=_lowercase			,			frame_sampling_rate=4		)

								A					       = hf_hub_download(repo_id='nateraw/video-demo'			,			filename='archery.mp4'			,			repo_type='dataset'		)
								A					       = video_classifier(_lowercase			,			top_k=2		)
								self.assertEqual(
								    nested_simplify(_lowercase			,			decimals=4		)			,			[{'score': 0.5_1_9_9, 'label': 'LABEL_0'}, {'score': 0.4_8_0_1, 'label': 'LABEL_1'}]			,			)

								A					       = video_classifier(
								    [
								        video_file_path,
								        video_file_path,
								    ]			,			top_k=2			,			)
								self.assertEqual(
								    nested_simplify(_lowercase			,			decimals=4		)			,			[
								        [{'score': 0.5_1_9_9, 'label': 'LABEL_0'}, {'score': 0.4_8_0_1, 'label': 'LABEL_1'}],
								        [{'score': 0.5_1_9_9, 'label': 'LABEL_0'}, {'score': 0.4_8_0_1, 'label': 'LABEL_1'}],
								    ]			,			)







						@require_tf
						def __a						(							self   :    Dict		):
								pass





