




"""simple docstring"""


import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch
     from torch import nn

     from transformers import (
         MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
         TimesformerForVideoClassification,
         TimesformerModel,
     )
     from transformers.models.timesformer.modeling_timesformer import TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from transformers import VideoMAEImageProcessor






class  lowercase:







 '''simple docstring'''





 def __init__(    self:		Any,				a_:		List[str],				a_:		Union[str, Any]=13,				a_:		List[str]=10,				a_:		Optional[Any]=3,				a_:		int=2,				a_:		Union[str, Any]=2,				a_:		Optional[int]=True,				a_:		int=True,				a_:		Tuple=32,				a_:		str=5,				a_:		str=4,				a_:		Tuple=37,				a_:		Optional[int]="gelu",				a_:		List[Any]=0.1,				a_:		Tuple=0.1,				a_:		Dict=10,				a_:		Any=0.02,				a_:		str="divided_space_time",				a_:		Tuple=None,				):




   '''simple docstring'''



   _snake_case							:		str								=  parent
   _snake_case							:		int								=  batch_size
   _snake_case							:		List[str]								=  image_size
   _snake_case							:		Any								=  num_channels
   _snake_case							:		int								=  patch_size
   _snake_case							:		Dict								=  num_frames
   _snake_case							:		Optional[Any]								=  is_training
   _snake_case							:		Optional[int]								=  use_labels
   _snake_case							:		List[str]								=  hidden_size
   _snake_case							:		List[Any]								=  num_hidden_layers
   _snake_case							:		Tuple								=  num_attention_heads
   _snake_case							:		Union[str, Any]								=  intermediate_size
   _snake_case							:		List[str]								=  hidden_act
   _snake_case							:		List[Any]								=  hidden_dropout_prob
   _snake_case							:		Optional[Any]								=  attention_probs_dropout_prob
   _snake_case							:		str								=  attention_type
   _snake_case							:		Optional[int]								=  initializer_range
   _snake_case							:		Dict								=  scope
   _snake_case							:		Union[str, Any]								=  num_labels

   # in TimeSformer, the number of spatial tokens equals num_frames * num_patches per frame + 1 CLS token
   _snake_case							:		Union[str, Any]								=  (image_size // patch_size) ** 2
   _snake_case							:		Tuple								=  (num_frames) * self.num_patches_per_frame + 1


 def 			UpperCamelCase_       (    self:		Any		):




   '''simple docstring'''



   _snake_case							:		Dict								=  floats_tensor(
       [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]		)

   _snake_case							:		Optional[Any]								=  None
   if self.use_labels:
     _snake_case							:		List[Any]								=  ids_tensor([self.batch_size],				self.num_labels		)

   _snake_case							:		int								=  self.get_config()

   return config, pixel_values, labels


 def 			UpperCamelCase_       (    self:		Any		):




   '''simple docstring'''



   _snake_case							:		str								=  TimesformerConfig(
       image_size=self.image_size,				patch_size=self.patch_size,				num_channels=self.num_channels,				num_frames=self.num_frames,				hidden_size=self.hidden_size,				num_hidden_layers=self.num_hidden_layers,				num_attention_heads=self.num_attention_heads,				intermediate_size=self.intermediate_size,				hidden_act=self.hidden_act,				hidden_dropout_prob=self.hidden_dropout_prob,				attention_probs_dropout_prob=self.attention_probs_dropout_prob,				initializer_range=self.initializer_range,				attention_type=self.attention_type,				)
   _snake_case							:		Union[str, Any]								=  self.num_labels
   return config


 def 			UpperCamelCase_       (    self:		Union[str, Any],				a_:		Optional[int],				a_:		List[str],				a_:		str		):




   '''simple docstring'''



   _snake_case							:		str								=  TimesformerModel(config=a_		)
   model.to(a_		)
   model.eval()
   _snake_case							:		str								=  model(a_		)
   self.parent.assertEqual(result.last_hidden_state.shape,				(self.batch_size, self.seq_length, self.hidden_size)		)


 def 			UpperCamelCase_       (    self:		Optional[Any],				a_:		str,				a_:		Union[str, Any],				a_:		List[str]		):




   '''simple docstring'''



   _snake_case							:		Optional[Any]								=  TimesformerForVideoClassification(a_		)
   model.to(a_		)
   model.eval()

   _snake_case							:		Optional[int]								=  model(a_		)

   # verify the logits shape
   _snake_case							:		str								=  torch.Size((self.batch_size, self.num_labels)		)
   self.parent.assertEqual(result.logits.shape,				a_		)


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




   '''simple docstring'''



   _snake_case							:		Any								=  self.prepare_config_and_inputs()
   _snake_case		, _snake_case		, _snake_case							:		Optional[Any]								=  config_and_inputs
   _snake_case							:		Dict								=  {"""pixel_values""": pixel_values}
   return config, inputs_dict







@require_torch
class  lowercase(					__a	,      __a	,      unittest.TestCase			):







 '''simple docstring'''





 lowercase__			 = (TimesformerModel, TimesformerForVideoClassification) if is_torch_available() else ()
 lowercase__			 = (
     {"feature-extraction": TimesformerModel, "video-classification": TimesformerForVideoClassification}
     if is_torch_available()
     else {}
 )

 lowercase__			 = False
 lowercase__			 = False
 lowercase__			 = False
 lowercase__			 = False


 def 			UpperCamelCase_       (    self:		Optional[Any]		):




   '''simple docstring'''



   _snake_case							:		Optional[Any]								=  TimesformerModelTester(self		)
   _snake_case							:		Dict								=  ConfigTester(
       self,				config_class=a_,				has_text_modality=a_,				hidden_size=37		)


 def 			UpperCamelCase_       (    self:		Optional[Any],				a_:		int,				a_:		Tuple,				a_:		Any=False		):




   '''simple docstring'''



   _snake_case							:		List[Any]								=  copy.deepcopy(a_		)

   if return_labels:
     if model_class in get_values(a_		):
       _snake_case							:		Dict								=  torch.zeros(
           self.model_tester.batch_size,				dtype=torch.long,				device=a_		)

   return inputs_dict


 def 			UpperCamelCase_       (    self:		List[Any]		):




   '''simple docstring'''



   self.config_tester.run_common_tests()


 @unittest.skip(reason="""TimeSformer does not use inputs_embeds"""		)
 def 			UpperCamelCase_       (    self:		Dict		):




   '''simple docstring'''



   pass


 def 			UpperCamelCase_       (    self:		Tuple		):




   '''simple docstring'''



   _snake_case		, _snake_case							:		Optional[Any]								=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:		Union[str, Any]								=  model_class(a_		)
     self.assertIsInstance(model.get_input_embeddings(),				(nn.Module)		)
     _snake_case							:		List[str]								=  model.get_output_embeddings()
     self.assertTrue(x is None or isinstance(a_,				nn.Linear		)		)


 def 			UpperCamelCase_       (    self:		List[str]		):




   '''simple docstring'''



   _snake_case		, _snake_case							:		Tuple								=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:		Any								=  model_class(a_		)
     _snake_case							:		int								=  inspect.signature(model.forward		)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     _snake_case							:		Union[str, Any]								=  [*signature.parameters.keys()]

     _snake_case							:		Optional[int]								=  ["""pixel_values"""]
     self.assertListEqual(arg_names[:1],				a_		)


 def 			UpperCamelCase_       (    self:		Tuple		):




   '''simple docstring'''



   _snake_case							:		List[Any]								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a_		)


 def 			UpperCamelCase_       (    self:		Dict		):




   '''simple docstring'''



   _snake_case							:		Any								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_video_classification(*a_		)


 @slow
 def 			UpperCamelCase_       (    self:		Optional[Any]		):




   '''simple docstring'''



   for model_name in TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case							:		Optional[Any]								=  TimesformerModel.from_pretrained(a_		)
     self.assertIsNotNone(a_		)


 def 			UpperCamelCase_       (    self:		Tuple		):




   '''simple docstring'''



   if not self.has_attentions:
     pass

   else:
     _snake_case		, _snake_case							:		List[str]								=  self.model_tester.prepare_config_and_inputs_for_common()
     _snake_case							:		str								=  True

     for model_class in self.all_model_classes:
       _snake_case							:		int								=  self.model_tester.seq_length
       _snake_case							:		int								=  self.model_tester.num_frames

       _snake_case							:		Dict								=  True
       _snake_case							:		Union[str, Any]								=  False
       _snake_case							:		Dict								=  True
       _snake_case							:		Optional[Any]								=  model_class(a_		)
       model.to(a_		)
       model.eval()
       with torch.no_grad():
         _snake_case							:		List[str]								=  model(**self._prepare_for_class(a_,				a_		)		)
       _snake_case							:		Any								=  outputs.attentions
       self.assertEqual(len(a_		),				self.model_tester.num_hidden_layers		)

       # check that output_attentions also work using config
       del inputs_dict["output_attentions"]
       _snake_case							:		Any								=  True
       _snake_case							:		int								=  model_class(a_		)
       model.to(a_		)
       model.eval()
       with torch.no_grad():
         _snake_case							:		str								=  model(**self._prepare_for_class(a_,				a_		)		)
       _snake_case							:		Optional[Any]								=  outputs.attentions
       self.assertEqual(len(a_		),				self.model_tester.num_hidden_layers		)

       # attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
       self.assertListEqual(
           list(attentions[0].shape[-3:]		),				[self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1],				)
       _snake_case							:		List[Any]								=  len(a_		)

       # Check attention is always last and order is fine
       _snake_case							:		str								=  True
       _snake_case							:		Any								=  True
       _snake_case							:		str								=  model_class(a_		)
       model.to(a_		)
       model.eval()
       with torch.no_grad():
         _snake_case							:		Any								=  model(**self._prepare_for_class(a_,				a_		)		)

       self.assertEqual(out_len + 1,				len(a_		)		)

       _snake_case							:		List[Any]								=  outputs.attentions

       self.assertEqual(len(a_		),				self.model_tester.num_hidden_layers		)

       # attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
       self.assertListEqual(
           list(self_attentions[0].shape[-3:]		),				[self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1],				)


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




   '''simple docstring'''



   def check_hidden_states_output(a_:		Optional[Any],				a_:		List[str],				a_:		str		):
     _snake_case							:		str								=  model_class(a_		)
     model.to(a_		)
     model.eval()

     with torch.no_grad():
       _snake_case							:		Optional[int]								=  model(**self._prepare_for_class(a_,				a_		)		)

     _snake_case							:		Optional[Any]								=  outputs.hidden_states
     _snake_case							:		Any								=  self.model_tester.num_hidden_layers + 1
     self.assertEqual(len(a_		),				a_		)

     _snake_case							:		Tuple								=  self.model_tester.seq_length

     self.assertListEqual(
         list(hidden_states[0].shape[-2:]		),				[seq_length, self.model_tester.hidden_size],				)

   _snake_case		, _snake_case							:		str								=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:		Any								=  True
     check_hidden_states_output(a_,				a_,				a_		)

     # check that output_hidden_states also work using config
     del inputs_dict["output_hidden_states"]
     _snake_case							:		List[Any]								=  True

     check_hidden_states_output(a_,				a_,				a_		)







def  UpperCAmelCase__	():
  """simple docstring"""

  _snake_case							:		str								=  hf_hub_download(
      repo_id="""hf-internal-testing/spaghetti-video"""		,       filename="""eating_spaghetti.npy"""		,       repo_type="""dataset"""   )
  _snake_case							:		List[Any]								=  np.load(snake_case__   )
  return list(snake_case__   )






@require_torch
@require_vision
class  lowercase(					unittest.TestCase			):







 '''simple docstring'''





 @cached_property
 def 			UpperCamelCase_       (    self:		List[str]		):




   '''simple docstring'''



   return (
       VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5],				image_std=[0.5, 0.5, 0.5]		)
       if is_vision_available()
       else None
   )


 @slow
 def 			UpperCamelCase_       (    self:		List[str]		):




   '''simple docstring'''



   _snake_case							:		str								=  TimesformerForVideoClassification.from_pretrained("""facebook/timesformer-base-finetuned-k400"""		).to(
       a_		)

   _snake_case							:		Tuple								=  self.default_image_processor
   _snake_case							:		Union[str, Any]								=  prepare_video()
   _snake_case							:		Union[str, Any]								=  image_processor(video[:8],				return_tensors="""pt"""		).to(a_		)

   # forward pass
   with torch.no_grad():
     _snake_case							:		int								=  model(**a_		)

   # verify the logits
   _snake_case							:		Optional[int]								=  torch.Size((1, 400)		)
   self.assertEqual(outputs.logits.shape,				a_		)

   _snake_case							:		List[Any]								=  torch.tensor([-0.3_016, -0.7_713, -0.4_205]		).to(a_		)

   self.assertTrue(torch.allclose(outputs.logits[0, :3],				a_,				atol=1E-4		)		)








"""simple docstring"""


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,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class  lowercase(					__a			):







 '''simple docstring'''





 lowercase__			 = 42
 lowercase__			 = 42







class  lowercase(					nn.Module			):







 '''simple docstring'''





 lowercase__			 = 42
 lowercase__			 = (16, 32, 96, 2_56)
 lowercase__			 = jnp.floataa


 def 			UpperCamelCase_       (    self:		Optional[int]		):




   '''simple docstring'''



   _snake_case							:		List[str]								=  nn.Conv(
       self.block_out_channels[0],				kernel_size=(3, 3),				padding=((1, 1), (1, 1)),				dtype=self.dtype,				)

   _snake_case							:		int								=  []
   for i in range(len(self.block_out_channels		) - 1		):
     _snake_case							:		int								=  self.block_out_channels[i]
     _snake_case							:		Tuple								=  self.block_out_channels[i + 1]
     _snake_case							:		Dict								=  nn.Conv(
         a_,				kernel_size=(3, 3),				padding=((1, 1), (1, 1)),				dtype=self.dtype,				)
     blocks.append(a_		)
     _snake_case							:		List[Any]								=  nn.Conv(
         a_,				kernel_size=(3, 3),				strides=(2, 2),				padding=((1, 1), (1, 1)),				dtype=self.dtype,				)
     blocks.append(a_		)
   _snake_case							:		Any								=  blocks

   _snake_case							:		Optional[Any]								=  nn.Conv(
       self.conditioning_embedding_channels,				kernel_size=(3, 3),				padding=((1, 1), (1, 1)),				kernel_init=nn.initializers.zeros_init(),				bias_init=nn.initializers.zeros_init(),				dtype=self.dtype,				)


 def __call__(    self:		Optional[Any],				a_:		Optional[Any]		):




   '''simple docstring'''



   _snake_case							:		int								=  self.conv_in(a_		)
   _snake_case							:		Optional[int]								=  nn.silu(a_		)

   for block in self.blocks:
     _snake_case							:		Tuple								=  block(a_		)
     _snake_case							:		int								=  nn.silu(a_		)

   _snake_case							:		Optional[int]								=  self.conv_out(a_		)

   return embedding







@flax_register_to_config
class  lowercase(					nn.Module	,      __a	,      __a			):







 '''simple docstring'''





 lowercase__			 = 32
 lowercase__			 = 4
 lowercase__			 = (
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 lowercase__			 = False
 lowercase__			 = (3_20, 6_40, 12_80, 12_80)
 lowercase__			 = 2
 lowercase__			 = 8
 lowercase__			 = None
 lowercase__			 = 12_80
 lowercase__			 = 0.0
 lowercase__			 = False
 lowercase__			 = jnp.floataa
 lowercase__			 = True
 lowercase__			 = 0
 lowercase__			 = "rgb"
 lowercase__			 = (16, 32, 96, 2_56)


 def 			UpperCamelCase_       (    self:		int,				a_:		jax.random.KeyArray		):




   '''simple docstring'''



   _snake_case							:		str								=  (1, self.in_channels, self.sample_size, self.sample_size)
   _snake_case							:		Optional[Any]								=  jnp.zeros(a_,				dtype=jnp.floataa		)
   _snake_case							:		List[str]								=  jnp.ones((1,),				dtype=jnp.intaa		)
   _snake_case							:		List[Any]								=  jnp.zeros((1, 1, self.cross_attention_dim),				dtype=jnp.floataa		)
   _snake_case							:		Any								=  (1, 3, self.sample_size * 8, self.sample_size * 8)
   _snake_case							:		Optional[int]								=  jnp.zeros(a_,				dtype=jnp.floataa		)

   _snake_case		, _snake_case							:		Tuple								=  jax.random.split(a_		)
   _snake_case							:		str								=  {"""params""": params_rng, """dropout""": dropout_rng}

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


 def 			UpperCamelCase_       (    self:		Any		):




   '''simple docstring'''



   _snake_case							:		Optional[int]								=  self.block_out_channels
   _snake_case							:		Optional[int]								=  block_out_channels[0] * 4

   # 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.
   _snake_case							:		int								=  self.num_attention_heads or self.attention_head_dim

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

   # time
   _snake_case							:		int								=  FlaxTimesteps(
       block_out_channels[0],				flip_sin_to_cos=self.flip_sin_to_cos,				freq_shift=self.config.freq_shift		)
   _snake_case							:		Any								=  FlaxTimestepEmbedding(a_,				dtype=self.dtype		)

   _snake_case							:		Optional[Any]								=  FlaxControlNetConditioningEmbedding(
       conditioning_embedding_channels=block_out_channels[0],				block_out_channels=self.conditioning_embedding_out_channels,				)

   _snake_case							:		List[str]								=  self.only_cross_attention
   if isinstance(a_,				a_		):
     _snake_case							:		List[Any]								=  (only_cross_attention,) * len(self.down_block_types		)

   if isinstance(a_,				a_		):
     _snake_case							:		Optional[Any]								=  (num_attention_heads,) * len(self.down_block_types		)

   # down
   _snake_case							:		List[str]								=  []
   _snake_case							:		Tuple								=  []

   _snake_case							:		int								=  block_out_channels[0]

   _snake_case							:		Optional[Any]								=  nn.Conv(
       a_,				kernel_size=(1, 1),				padding="""VALID""",				kernel_init=nn.initializers.zeros_init(),				bias_init=nn.initializers.zeros_init(),				dtype=self.dtype,				)
   controlnet_down_blocks.append(a_		)

   for i, down_block_type in enumerate(self.down_block_types		):
     _snake_case							:		List[Any]								=  output_channel
     _snake_case							:		Any								=  block_out_channels[i]
     _snake_case							:		List[str]								=  i == len(a_		) - 1

     if down_block_type == "CrossAttnDownBlock2D":
       _snake_case							:		Optional[int]								=  FlaxCrossAttnDownBlockaD(
           in_channels=a_,				out_channels=a_,				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],				dtype=self.dtype,				)
     else:
       _snake_case							:		List[Any]								=  FlaxDownBlockaD(
           in_channels=a_,				out_channels=a_,				dropout=self.dropout,				num_layers=self.layers_per_block,				add_downsample=not is_final_block,				dtype=self.dtype,				)

     down_blocks.append(a_		)

     for _ in range(self.layers_per_block		):
       _snake_case							:		List[Any]								=  nn.Conv(
           a_,				kernel_size=(1, 1),				padding="""VALID""",				kernel_init=nn.initializers.zeros_init(),				bias_init=nn.initializers.zeros_init(),				dtype=self.dtype,				)
       controlnet_down_blocks.append(a_		)

     if not is_final_block:
       _snake_case							:		List[Any]								=  nn.Conv(
           a_,				kernel_size=(1, 1),				padding="""VALID""",				kernel_init=nn.initializers.zeros_init(),				bias_init=nn.initializers.zeros_init(),				dtype=self.dtype,				)
       controlnet_down_blocks.append(a_		)

   _snake_case							:		str								=  down_blocks
   _snake_case							:		Union[str, Any]								=  controlnet_down_blocks

   # mid
   _snake_case							:		Tuple								=  block_out_channels[-1]
   _snake_case							:		Optional[int]								=  FlaxUNetMidBlockaDCrossAttn(
       in_channels=a_,				dropout=self.dropout,				num_attention_heads=num_attention_heads[-1],				use_linear_projection=self.use_linear_projection,				dtype=self.dtype,				)

   _snake_case							:		Tuple								=  nn.Conv(
       a_,				kernel_size=(1, 1),				padding="""VALID""",				kernel_init=nn.initializers.zeros_init(),				bias_init=nn.initializers.zeros_init(),				dtype=self.dtype,				)


 def __call__(    self:		str,				a_:		Any,				a_:		Tuple,				a_:		Any,				a_:		int,				a_:		float = 1.0,				a_:		bool = True,				a_:		bool = False,				):




   '''simple docstring'''



   _snake_case							:		Dict								=  self.controlnet_conditioning_channel_order
   if channel_order == "bgr":
     _snake_case							:		List[Any]								=  jnp.flip(a_,				axis=1		)

   # 1. time
   if not isinstance(a_,				jnp.ndarray		):
     _snake_case							:		Any								=  jnp.array([timesteps],				dtype=jnp.intaa		)
   elif isinstance(a_,				jnp.ndarray		) and len(timesteps.shape		) == 0:
     _snake_case							:		Union[str, Any]								=  timesteps.astype(dtype=jnp.floataa		)
     _snake_case							:		List[str]								=  jnp.expand_dims(a_,				0		)

   _snake_case							:		List[str]								=  self.time_proj(a_		)
   _snake_case							:		str								=  self.time_embedding(a_		)

   # 2. pre-process
   _snake_case							:		List[str]								=  jnp.transpose(a_,				(0, 2, 3, 1)		)
   _snake_case							:		List[Any]								=  self.conv_in(a_		)

   _snake_case							:		Union[str, Any]								=  jnp.transpose(a_,				(0, 2, 3, 1)		)
   _snake_case							:		Any								=  self.controlnet_cond_embedding(a_		)
   sample += controlnet_cond

   # 3. down
   _snake_case							:		List[str]								=  (sample,)
   for down_block in self.down_blocks:
     if isinstance(a_,				a_		):
       _snake_case		, _snake_case							:		Optional[Any]								=  down_block(a_,				a_,				a_,				deterministic=not train		)
     else:
       _snake_case		, _snake_case							:		Dict								=  down_block(a_,				a_,				deterministic=not train		)
     down_block_res_samples += res_samples

   # 4. mid
   _snake_case							:		Dict								=  self.mid_block(a_,				a_,				a_,				deterministic=not train		)

   # 5. contronet blocks
   _snake_case							:		Tuple								=  ()
   for down_block_res_sample, controlnet_block in zip(a_,				self.controlnet_down_blocks		):
     _snake_case							:		Any								=  controlnet_block(a_		)
     controlnet_down_block_res_samples += (down_block_res_sample,)

   _snake_case							:		List[Any]								=  controlnet_down_block_res_samples

   _snake_case							:		int								=  self.controlnet_mid_block(a_		)

   # 6. scaling
   _snake_case							:		int								=  [sample * conditioning_scale for sample in down_block_res_samples]
   mid_block_res_sample *= conditioning_scale

   if not return_dict:
     return (down_block_res_samples, mid_block_res_sample)

   return FlaxControlNetOutput(
       down_block_res_samples=a_,				mid_block_res_sample=a_		)


