
from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)




@flax.struct.dataclass
class 				_lowerCamelCase (      UpperCamelCase_ ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  42




@flax_register_to_config
class 				_lowerCamelCase (      nn.Module	,							UpperCamelCase_	,							UpperCamelCase_ ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  32
     SCREAMING_SNAKE_CASE_		   =  4
     SCREAMING_SNAKE_CASE_		   =  4
     SCREAMING_SNAKE_CASE_		   =  (
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "DownBlock2D",
     )
     SCREAMING_SNAKE_CASE_		   =  ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
     SCREAMING_SNAKE_CASE_		   =  False
     SCREAMING_SNAKE_CASE_		   =  (320, 640, 1280, 1280)
     SCREAMING_SNAKE_CASE_		   =  2
     SCREAMING_SNAKE_CASE_		   =  8
     SCREAMING_SNAKE_CASE_		   =  None
     SCREAMING_SNAKE_CASE_		   =  1280
     SCREAMING_SNAKE_CASE_		   =  0.0
     SCREAMING_SNAKE_CASE_		   =  False
     SCREAMING_SNAKE_CASE_		   =  jnp.floataa
     SCREAMING_SNAKE_CASE_		   =  True
     SCREAMING_SNAKE_CASE_		   =  0
     SCREAMING_SNAKE_CASE_		   =  False






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   FrozenDict:




           """simple docstring"""







           UpperCamelCase__    :   List[Any]					=				(1, self.in_channels, self.sample_size, self.sample_size)
           UpperCamelCase__    :   Tuple					=				jnp.zeros(__SCREAMING_SNAKE_CASE						,						dtype=jnp.floataa  )
           UpperCamelCase__    :   Optional[int]					=				jnp.ones((1,)						,						dtype=jnp.intaa  )
           UpperCamelCase__    :   List[str]					=				jnp.zeros((1, 1, self.cross_attention_dim)						,						dtype=jnp.floataa  )

           UpperCamelCase__				,UpperCamelCase__    :   Union[str, Any]					=				jax.random.split(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Optional[Any]					=				{'''params''': params_rng, '''dropout''': dropout_rng}

           return self.init(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )["params"]






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   Dict					=				self.block_out_channels
           UpperCamelCase__    :   Any					=				block_out_channels[0] * 4

           if self.num_attention_heads is not None:
                 raise ValueError(
                     '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'''  )

           # If `num_attention_heads` is not defined (which is the case for most models)
           # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
           # The reason for this behavior is to correct for incorrectly named variables that were introduced
           # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
           # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
           # which is why we correct for the naming here.
           UpperCamelCase__    :   List[Any]					=				self.num_attention_heads or self.attention_head_dim

           # input
           UpperCamelCase__    :   str					=				nn.Conv(
               block_out_channels[0]						,						kernel_size=(3, 3)						,						strides=(1, 1)						,						padding=((1, 1), (1, 1))						,						dtype=self.dtype						,						)

           # time
           UpperCamelCase__    :   List[Any]					=				FlaxTimesteps(
               block_out_channels[0]						,						flip_sin_to_cos=self.flip_sin_to_cos						,						freq_shift=self.config.freq_shift  )
           UpperCamelCase__    :   str					=				FlaxTimestepEmbedding(__SCREAMING_SNAKE_CASE						,						dtype=self.dtype  )

           UpperCamelCase__    :   Any					=				self.only_cross_attention
           if isinstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  ):
                 UpperCamelCase__    :   str					=				(only_cross_attention,) * len(self.down_block_types  )

           if isinstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  ):
                 UpperCamelCase__    :   str					=				(num_attention_heads,) * len(self.down_block_types  )

           # down
           UpperCamelCase__    :   Optional[int]					=				[]
           UpperCamelCase__    :   Any					=				block_out_channels[0]
           for i, down_block_type in enumerate(self.down_block_types  ):
                 UpperCamelCase__    :   Dict					=				output_channel
                 UpperCamelCase__    :   int					=				block_out_channels[i]
                 UpperCamelCase__    :   Dict					=				i == len(__SCREAMING_SNAKE_CASE  ) - 1

                 if down_block_type == "CrossAttnDownBlock2D":
                       UpperCamelCase__    :   List[str]					=				FlaxCrossAttnDownBlockaD(
                           in_channels=__SCREAMING_SNAKE_CASE						,						out_channels=__SCREAMING_SNAKE_CASE						,						dropout=self.dropout						,						num_layers=self.layers_per_block						,						num_attention_heads=num_attention_heads[i]						,						add_downsample=not is_final_block						,						use_linear_projection=self.use_linear_projection						,						only_cross_attention=only_cross_attention[i]						,						use_memory_efficient_attention=self.use_memory_efficient_attention						,						dtype=self.dtype						,						)
                 else:
                       UpperCamelCase__    :   Union[str, Any]					=				FlaxDownBlockaD(
                           in_channels=__SCREAMING_SNAKE_CASE						,						out_channels=__SCREAMING_SNAKE_CASE						,						dropout=self.dropout						,						num_layers=self.layers_per_block						,						add_downsample=not is_final_block						,						dtype=self.dtype						,						)

                 down_blocks.append(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   str					=				down_blocks

           # mid
           UpperCamelCase__    :   str					=				FlaxUNetMidBlockaDCrossAttn(
               in_channels=block_out_channels[-1]						,						dropout=self.dropout						,						num_attention_heads=num_attention_heads[-1]						,						use_linear_projection=self.use_linear_projection						,						use_memory_efficient_attention=self.use_memory_efficient_attention						,						dtype=self.dtype						,						)

           # up
           UpperCamelCase__    :   Tuple					=				[]
           UpperCamelCase__    :   int					=				list(reversed(__SCREAMING_SNAKE_CASE  )  )
           UpperCamelCase__    :   str					=				list(reversed(__SCREAMING_SNAKE_CASE  )  )
           UpperCamelCase__    :   Tuple					=				list(reversed(__SCREAMING_SNAKE_CASE  )  )
           UpperCamelCase__    :   Optional[Any]					=				reversed_block_out_channels[0]
           for i, up_block_type in enumerate(self.up_block_types  ):
                 UpperCamelCase__    :   Optional[Any]					=				output_channel
                 UpperCamelCase__    :   Optional[int]					=				reversed_block_out_channels[i]
                 UpperCamelCase__    :   List[Any]					=				reversed_block_out_channels[min(i + 1						,						len(__SCREAMING_SNAKE_CASE  ) - 1  )]

                 UpperCamelCase__    :   Optional[int]					=				i == len(__SCREAMING_SNAKE_CASE  ) - 1

                 if up_block_type == "CrossAttnUpBlock2D":
                       UpperCamelCase__    :   Optional[Any]					=				FlaxCrossAttnUpBlockaD(
                           in_channels=__SCREAMING_SNAKE_CASE						,						out_channels=__SCREAMING_SNAKE_CASE						,						prev_output_channel=__SCREAMING_SNAKE_CASE						,						num_layers=self.layers_per_block + 1						,						num_attention_heads=reversed_num_attention_heads[i]						,						add_upsample=not is_final_block						,						dropout=self.dropout						,						use_linear_projection=self.use_linear_projection						,						only_cross_attention=only_cross_attention[i]						,						use_memory_efficient_attention=self.use_memory_efficient_attention						,						dtype=self.dtype						,						)
                 else:
                       UpperCamelCase__    :   List[str]					=				FlaxUpBlockaD(
                           in_channels=__SCREAMING_SNAKE_CASE						,						out_channels=__SCREAMING_SNAKE_CASE						,						prev_output_channel=__SCREAMING_SNAKE_CASE						,						num_layers=self.layers_per_block + 1						,						add_upsample=not is_final_block						,						dropout=self.dropout						,						dtype=self.dtype						,						)

                 up_blocks.append(__SCREAMING_SNAKE_CASE  )
                 UpperCamelCase__    :   Optional[int]					=				output_channel
           UpperCamelCase__    :   Tuple					=				up_blocks

           # out
           UpperCamelCase__    :   Union[str, Any]					=				nn.GroupNorm(num_groups=3_2						,						epsilon=1e-5  )
           UpperCamelCase__    :   Any					=				nn.Conv(
               self.out_channels						,						kernel_size=(3, 3)						,						strides=(1, 1)						,						padding=((1, 1), (1, 1))						,						dtype=self.dtype						,						)






     def __call__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=None						,						__SCREAMING_SNAKE_CASE=None						,						__SCREAMING_SNAKE_CASE = True						,						__SCREAMING_SNAKE_CASE = False						,						)  ->   Union[FlaxUNetaDConditionOutput, Tuple]:




           """simple docstring"""







           if not isinstance(__SCREAMING_SNAKE_CASE						,						jnp.ndarray  ):
                 UpperCamelCase__    :   List[Any]					=				jnp.array([timesteps]						,						dtype=jnp.intaa  )
           elif isinstance(__SCREAMING_SNAKE_CASE						,						jnp.ndarray  ) and len(timesteps.shape  ) == 0:
                 UpperCamelCase__    :   Optional[Any]					=				timesteps.astype(dtype=jnp.floataa  )
                 UpperCamelCase__    :   str					=				jnp.expand_dims(__SCREAMING_SNAKE_CASE						,						0  )

           UpperCamelCase__    :   Optional[int]					=				self.time_proj(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   List[str]					=				self.time_embedding(__SCREAMING_SNAKE_CASE  )

           # 2. pre-process
           UpperCamelCase__    :   Union[str, Any]					=				jnp.transpose(__SCREAMING_SNAKE_CASE						,						(0, 2, 3, 1)  )
           UpperCamelCase__    :   Union[str, Any]					=				self.conv_in(__SCREAMING_SNAKE_CASE  )

           # 3. down
           UpperCamelCase__    :   int					=				(sample,)
           for down_block in self.down_blocks:
                 if isinstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  ):
                       UpperCamelCase__				,UpperCamelCase__    :   Any					=				down_block(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						deterministic=not train  )
                 else:
                       UpperCamelCase__				,UpperCamelCase__    :   Tuple					=				down_block(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						deterministic=not train  )
                 down_block_res_samples += res_samples

           if down_block_additional_residuals is not None:
                 UpperCamelCase__    :   str					=				()

                 for down_block_res_sample, down_block_additional_residual in zip(
                     __SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  ):
                       down_block_res_sample += down_block_additional_residual
                       new_down_block_res_samples += (down_block_res_sample,)

                 UpperCamelCase__    :   Optional[Any]					=				new_down_block_res_samples

           # 4. mid
           UpperCamelCase__    :   Dict					=				self.mid_block(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						deterministic=not train  )

           if mid_block_additional_residual is not None:
                 sample += mid_block_additional_residual

           # 5. up
           for up_block in self.up_blocks:
                 UpperCamelCase__    :   List[str]					=				down_block_res_samples[-(self.layers_per_block + 1) :]
                 UpperCamelCase__    :   Optional[Any]					=				down_block_res_samples[: -(self.layers_per_block + 1)]
                 if isinstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  ):
                       UpperCamelCase__    :   List[str]					=				up_block(
                           __SCREAMING_SNAKE_CASE						,						temb=__SCREAMING_SNAKE_CASE						,						encoder_hidden_states=__SCREAMING_SNAKE_CASE						,						res_hidden_states_tuple=__SCREAMING_SNAKE_CASE						,						deterministic=not train						,						)
                 else:
                       UpperCamelCase__    :   Union[str, Any]					=				up_block(__SCREAMING_SNAKE_CASE						,						temb=__SCREAMING_SNAKE_CASE						,						res_hidden_states_tuple=__SCREAMING_SNAKE_CASE						,						deterministic=not train  )

        # 6. post-process
           UpperCamelCase__    :   List[Any]					=				self.conv_norm_out(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Optional[Any]					=				nn.silu(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   List[str]					=				self.conv_out(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Union[str, Any]					=				jnp.transpose(__SCREAMING_SNAKE_CASE						,						(0, 3, 1, 2)  )

           if not return_dict:
                 return (sample,)

           return FlaxUNetaDConditionOutput(sample=__SCREAMING_SNAKE_CASE  )




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(UpperCamelCase_ )
class 				_lowerCamelCase (      UpperCamelCase_ ):





     """simple docstring"""







     def __init__(							self						,						*__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Dict:




           """simple docstring"""







           super().__init__(*__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )
           requires_backends(self						,						'''decord'''  )
           self.check_model_type(__SCREAMING_SNAKE_CASE  )






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE=None						,						__SCREAMING_SNAKE_CASE=None						,						__SCREAMING_SNAKE_CASE=None  )  ->   Optional[Any]:




           """simple docstring"""







           UpperCamelCase__    :   List[str]					=				{}
           if frame_sampling_rate is not None:
                 UpperCamelCase__    :   Tuple					=				frame_sampling_rate
           if num_frames is not None:
                 UpperCamelCase__    :   str					=				num_frames

           UpperCamelCase__    :   List[str]					=				{}
           if top_k is not None:
                 UpperCamelCase__    :   List[Any]					=				top_k
           return preprocess_params, {}, postprocess_params






     def __call__(							self						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Tuple:




           """simple docstring"""







           return super().__call__(__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=None						,						__SCREAMING_SNAKE_CASE=1  )  ->   Optional[Any]:




           """simple docstring"""







           if num_frames is None:
                 UpperCamelCase__    :   Optional[Any]					=				self.model.config.num_frames

           if video.startswith('''http://'''  ) or video.startswith('''https://'''  ):
                 UpperCamelCase__    :   str					=				BytesIO(requests.get(__SCREAMING_SNAKE_CASE  ).content  )

           UpperCamelCase__    :   Tuple					=				VideoReader(__SCREAMING_SNAKE_CASE  )
           videoreader.seek(0  )

           UpperCamelCase__    :   Optional[int]					=				0
           UpperCamelCase__    :   str					=				num_frames * frame_sampling_rate - 1
           UpperCamelCase__    :   List[Any]					=				np.linspace(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						num=__SCREAMING_SNAKE_CASE						,						dtype=np.intaa  )

           UpperCamelCase__    :   str					=				videoreader.get_batch(__SCREAMING_SNAKE_CASE  ).asnumpy()
           UpperCamelCase__    :   Tuple					=				list(__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Dict					=				self.image_processor(__SCREAMING_SNAKE_CASE						,						return_tensors=self.framework  )
           return model_inputs






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE  )  ->   Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__    :   Tuple					=				self.model(**__SCREAMING_SNAKE_CASE  )
           return model_outputs






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=5  )  ->   str:




           """simple docstring"""







           if top_k > self.model.config.num_labels:
                 UpperCamelCase__    :   Dict					=				self.model.config.num_labels

           if self.framework == "pt":
                 UpperCamelCase__    :   Any					=				model_outputs.logits.softmax(-1  )[0]
                 UpperCamelCase__				,UpperCamelCase__    :   List[str]					=				probs.topk(__SCREAMING_SNAKE_CASE  )
           else:
                 raise ValueError(F'''Unsupported framework: {self.framework}'''  )

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


