

"""simple docstring"""




import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_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 ASTForAudioClassification, ASTModel
      from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
          AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      )


if is_torchaudio_available():
      import torchaudio

      from transformers import ASTFeatureExtractor







class   __lowercase      :

  def __init__(	self		:      List[str]       ,A		:      Optional[Any]       ,A		:      Optional[int]=13       ,A		:      Tuple=2       ,A		:      int=24       ,A		:      int=16       ,A		:      Dict=True       ,A		:      Tuple=True       ,A		:      List[str]=32       ,A		:      Tuple=5       ,A		:      Optional[int]=4       ,A		:      List[str]=37       ,A		:      List[Any]="gelu"       ,A		:      int=0.1       ,A		:      Tuple=0.1       ,A		:      Optional[int]=10       ,A		:      Optional[Any]=0.0_2       ,A		:      Optional[int]=None       ,A		:      Optional[int]=2       ,A		:      Optional[Any]=2       ,):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       parent
   UpperCAmelCase__       :	Union[str, Any]        =       batch_size
   UpperCAmelCase__       :	int        =       patch_size
   UpperCAmelCase__       :	Any        =       max_length
   UpperCAmelCase__       :	int        =       num_mel_bins
   UpperCAmelCase__       :	Tuple        =       is_training
   UpperCAmelCase__       :	Tuple        =       use_labels
   UpperCAmelCase__       :	Union[str, Any]        =       hidden_size
   UpperCAmelCase__       :	str        =       num_hidden_layers
   UpperCAmelCase__       :	List[str]        =       num_attention_heads
   UpperCAmelCase__       :	str        =       intermediate_size
   UpperCAmelCase__       :	List[Any]        =       hidden_act
   UpperCAmelCase__       :	Union[str, Any]        =       hidden_dropout_prob
   UpperCAmelCase__       :	Any        =       attention_probs_dropout_prob
   UpperCAmelCase__       :	Optional[int]        =       type_sequence_label_size
   UpperCAmelCase__       :	Tuple        =       initializer_range
   UpperCAmelCase__       :	List[str]        =       scope
   UpperCAmelCase__       :	Dict        =       frequency_stride
   UpperCAmelCase__       :	Tuple        =       time_stride

   # in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
   UpperCAmelCase__       :	Tuple        =       (self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
   UpperCAmelCase__       :	Optional[Any]        =       (self.max_length - self.patch_size) // self.time_stride + 1
   UpperCAmelCase__       :	List[Any]        =       frequency_out_dimension * time_out_dimension
   UpperCAmelCase__       :	Any        =       num_patches + 2

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]     )

   UpperCAmelCase__       :	Tuple        =       None
   if self.use_labels:
    UpperCAmelCase__       :	List[str]        =       ids_tensor([self.batch_size]       ,self.type_sequence_label_size     )

   UpperCAmelCase__       :	Union[str, Any]        =       self.get_config()

   return config, input_values, labels

  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   return ASTConfig(
       patch_size=self.patch_size       ,max_length=self.max_length       ,num_mel_bins=self.num_mel_bins       ,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       ,is_decoder=__SCREAMING_SNAKE_CASE       ,initializer_range=self.initializer_range       ,frequency_stride=self.frequency_stride       ,time_stride=self.time_stride       ,)

  def  __lowercase							(	self		:      Dict       ,A		:      Optional[Any]       ,A		:      Union[str, Any]       ,A		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       ASTModel(config=__SCREAMING_SNAKE_CASE     )
   model.to(__SCREAMING_SNAKE_CASE     )
   model.eval()
   UpperCAmelCase__       :	List[Any]        =       model(__SCREAMING_SNAKE_CASE     )
   self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size)     )




  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Union[str, Any]        =       self.prepare_config_and_inputs()
   (
       (
       UpperCAmelCase__     
   ) ,   (
       UpperCAmelCase__     
   ) ,   (
       UpperCAmelCase__     
   ) ,        
   )       :	int        =       config_and_inputs
   UpperCAmelCase__       :	Tuple        =       {"""input_values""": input_values}
   return config, inputs_dict







@require_torch
class   __lowercase      (     _lowercase						,      _lowercase						,      unittest.TestCase       ):
  snake_case_												=							(
      (
          ASTModel,
          ASTForAudioClassification,
      )
      if is_torch_available()
      else ()
  )
  snake_case_												=							(
      {'audio-classification': ASTForAudioClassification, 'feature-extraction': ASTModel}
      if is_torch_available()
      else {}
  )
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False

  def  __lowercase							(	self		:      Tuple       ,A		:      Tuple       ,A		:      Union[str, Any]       ,A		:      Optional[int]       ,A		:      int       ,A		:      Tuple     ):






   '''simple docstring'''



   if pipeline_test_casse_name == "AudioClassificationPipelineTests":
    return True

   return False

  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       ASTModelTester(self     )
   UpperCAmelCase__       :	Any        =       ConfigTester(self       ,config_class=__SCREAMING_SNAKE_CASE       ,has_text_modality=__SCREAMING_SNAKE_CASE       ,hidden_size=37     )

  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   self.config_tester.run_common_tests()

  @unittest.skip(reason="""AST does not use inputs_embeds"""     )
  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   pass

  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   UpperCAmelCase__ ,   UpperCAmelCase__       :	List[str]        =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	str        =       model_class(__SCREAMING_SNAKE_CASE     )
    self.assertIsInstance(model.get_input_embeddings()       ,(nn.Module)     )
    UpperCAmelCase__       :	int        =       model.get_output_embeddings()
    self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE       ,nn.Linear     )     )

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






   '''simple docstring'''
   UpperCAmelCase__ ,   UpperCAmelCase__       :	Union[str, Any]        =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	Optional[Any]        =       model_class(__SCREAMING_SNAKE_CASE     )
    UpperCAmelCase__       :	List[str]        =       inspect.signature(model.forward     )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    UpperCAmelCase__       :	Dict        =       [*signature.parameters.keys()]

    UpperCAmelCase__       :	Tuple        =       ["""input_values"""]
    self.assertListEqual(arg_names[:1]       ,__SCREAMING_SNAKE_CASE     )

  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE     )




  @slow
  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''



   for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    UpperCAmelCase__       :	List[Any]        =       ASTModel.from_pretrained(__SCREAMING_SNAKE_CASE     )
    self.assertIsNotNone(__SCREAMING_SNAKE_CASE     )







def 	lowerCAmelCase    (       ):
 '''simple docstring'''





 UpperCAmelCase__       :	Optional[int]        =       hf_hub_download(
     repo_id="""nielsr/audio-spectogram-transformer-checkpoint"""				,     filename="""sample_audio.flac"""				,     repo_type="""dataset"""   )

 UpperCAmelCase__ ,   UpperCAmelCase__       :	List[Any]        =       torchaudio.load(__UpperCamelCase   )

 return audio, sampling_rate







@require_torch
@require_torchaudio
class   __lowercase      (     unittest.TestCase       ):

  @cached_property
  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   return (
       ASTFeatureExtractor.from_pretrained("""MIT/ast-finetuned-audioset-10-10-0.4593"""     )
       if is_torchaudio_available()
       else None
   )




  @slow
  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       self.default_feature_extractor
   UpperCAmelCase__       :	Tuple        =       ASTForAudioClassification.from_pretrained("""MIT/ast-finetuned-audioset-10-10-0.4593"""     ).to(__SCREAMING_SNAKE_CASE     )

   UpperCAmelCase__       :	Any        =       self.default_feature_extractor
   UpperCAmelCase__ ,   UpperCAmelCase__       :	Any        =       prepare_audio()
   UpperCAmelCase__       :	List[Any]        =       audio.squeeze().numpy()
   UpperCAmelCase__       :	int        =       feature_extractor(__SCREAMING_SNAKE_CASE       ,sampling_rate=__SCREAMING_SNAKE_CASE       ,return_tensors="""pt"""     ).to(__SCREAMING_SNAKE_CASE     )

   # forward pass
   with torch.no_grad():
    UpperCAmelCase__       :	Dict        =       model(**__SCREAMING_SNAKE_CASE     )

   # verify the logits
   UpperCAmelCase__       :	Optional[int]        =       torch.Size((1, 527)     )
   self.assertEqual(outputs.logits.shape       ,__SCREAMING_SNAKE_CASE     )

   UpperCAmelCase__       :	List[str]        =       torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]     ).to(__SCREAMING_SNAKE_CASE     )

   self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,__SCREAMING_SNAKE_CASE       ,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(_lowercase       ):
     __snake_case:	jnp.ndarray
     __snake_case:	jnp.ndarray







class     lowercase(nn.Module       ):
     __snake_case:	int
     __snake_case:	Tuple[int]						=     (16, 32, 96, 256)
     __snake_case:	jnp.dtype						=     jnp.floataa



     def 			lowercase__  (							self   )		->     Dict:







         """simple docstring"""



         a__		=			nn.Conv(
             self.block_out_channels[0]     ,				kernel_size=(3, 3)     ,				padding=((1, 1), (1, 1))     ,				dtype=self.dtype     ,				)

         a__		=			[]
         for i in range(len(self.block_out_channels   ) - 1   ):
             a__		=			self.block_out_channels[i]
             a__		=			self.block_out_channels[i + 1]
             a__		=			nn.Conv(
                 __SCREAMING_SNAKE_CASE     ,				kernel_size=(3, 3)     ,				padding=((1, 1), (1, 1))     ,				dtype=self.dtype     ,				)
             blocks.append(__SCREAMING_SNAKE_CASE   )
             a__		=			nn.Conv(
                 __SCREAMING_SNAKE_CASE     ,				kernel_size=(3, 3)     ,				strides=(2, 2)     ,				padding=((1, 1), (1, 1))     ,				dtype=self.dtype     ,				)
             blocks.append(__SCREAMING_SNAKE_CASE   )
         a__		=			blocks

         a__		=			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     ,				__SCREAMING_SNAKE_CASE   )		->     Any:







         """simple docstring"""



         a__		=			self.conv_in(__SCREAMING_SNAKE_CASE   )
         a__		=			nn.silu(__SCREAMING_SNAKE_CASE   )

         for block in self.blocks:
             a__		=			block(__SCREAMING_SNAKE_CASE   )
             a__		=			nn.silu(__SCREAMING_SNAKE_CASE   )

         a__		=			self.conv_out(__SCREAMING_SNAKE_CASE   )

         return embedding







@flax_register_to_config
class     lowercase(nn.Module    ,						_lowercase    ,						_lowercase       ):
     __snake_case:	int						=     32
     __snake_case:	int						=     4
     __snake_case:	Tuple[str]						=     (
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "DownBlock2D",
     )
     __snake_case:	Union[bool, Tuple[bool]]						=     False
     __snake_case:	Tuple[int]						=     (320, 640, 1280, 1280)
     __snake_case:	int						=     2
     __snake_case:	Union[int, Tuple[int]]						=     8
     __snake_case:	Optional[Union[int, Tuple[int]]]						=     None
     __snake_case:	int						=     1280
     __snake_case:	float						=     0.0
     __snake_case:	bool						=     False
     __snake_case:	jnp.dtype						=     jnp.floataa
     __snake_case:	bool						=     True
     __snake_case:	int						=     0
     __snake_case:	str						=     "rgb"
     __snake_case:	Tuple[int]						=     (16, 32, 96, 256)



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE   )		->     FrozenDict:







         """simple docstring"""



         a__		=			(1, self.in_channels, self.sample_size, self.sample_size)
         a__		=			jnp.zeros(__SCREAMING_SNAKE_CASE     ,				dtype=jnp.floataa   )
         a__		=			jnp.ones((1,)     ,				dtype=jnp.intaa   )
         a__		=			jnp.zeros((1, 1, self.cross_attention_dim)     ,				dtype=jnp.floataa   )
         a__		=			(1, 3, self.sample_size * 8, self.sample_size * 8)
         a__		=			jnp.zeros(__SCREAMING_SNAKE_CASE     ,				dtype=jnp.floataa   )

         a__	,     a__		=			jax.random.split(__SCREAMING_SNAKE_CASE   )
         a__		=			{'params': params_rng, 'dropout': dropout_rng}

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



     def 			lowercase__  (							self   )		->     str:







         """simple docstring"""



         a__		=			self.block_out_channels
         a__		=			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.
         a__		=			self.num_attention_heads or self.attention_head_dim

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

         # time
         a__		=			FlaxTimesteps(
             block_out_channels[0]     ,				flip_sin_to_cos=self.flip_sin_to_cos     ,				freq_shift=self.config.freq_shift   )
         a__		=			FlaxTimestepEmbedding(__SCREAMING_SNAKE_CASE     ,				dtype=self.dtype   )

         a__		=			FlaxControlNetConditioningEmbedding(
             conditioning_embedding_channels=block_out_channels[0]     ,				block_out_channels=self.conditioning_embedding_out_channels     ,				)

         a__		=			self.only_cross_attention
         if isinstance(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   ):
             a__		=			(only_cross_attention,) * len(self.down_block_types   )

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

         # down
         a__		=			[]
         a__		=			[]

         a__		=			block_out_channels[0]

         a__		=			nn.Conv(
             __SCREAMING_SNAKE_CASE     ,				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(__SCREAMING_SNAKE_CASE   )

         for i, down_block_type in enumerate(self.down_block_types   ):
             a__		=			output_channel
             a__		=			block_out_channels[i]
             a__		=			i == len(__SCREAMING_SNAKE_CASE   ) - 1

             if down_block_type == "CrossAttnDownBlock2D":
                 a__		=			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]     ,				dtype=self.dtype     ,				)
             else:
                 a__		=			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   )

             for _ in range(self.layers_per_block   ):
                 a__		=			nn.Conv(
                     __SCREAMING_SNAKE_CASE     ,				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(__SCREAMING_SNAKE_CASE   )

             if not is_final_block:
                 a__		=			nn.Conv(
                     __SCREAMING_SNAKE_CASE     ,				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(__SCREAMING_SNAKE_CASE   )

         a__		=			down_blocks
         a__		=			controlnet_down_blocks

         # mid
         a__		=			block_out_channels[-1]
         a__		=			FlaxUNetMidBlockaDCrossAttn(
             in_channels=__SCREAMING_SNAKE_CASE     ,				dropout=self.dropout     ,				num_attention_heads=num_attention_heads[-1]     ,				use_linear_projection=self.use_linear_projection     ,				dtype=self.dtype     ,				)

         a__		=			nn.Conv(
             __SCREAMING_SNAKE_CASE     ,				kernel_size=(1, 1)     ,				padding='VALID'     ,				kernel_init=nn.initializers.zeros_init()     ,				bias_init=nn.initializers.zeros_init()     ,				dtype=self.dtype     ,				)




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







         """simple docstring"""



         a__		=			self.controlnet_conditioning_channel_order
         if channel_order == "bgr":
             a__		=			jnp.flip(__SCREAMING_SNAKE_CASE     ,				axis=1   )

         # 1. time
         if not isinstance(__SCREAMING_SNAKE_CASE     ,				jnp.ndarray   ):
             a__		=			jnp.array([timesteps]     ,				dtype=jnp.intaa   )
         elif isinstance(__SCREAMING_SNAKE_CASE     ,				jnp.ndarray   ) and len(timesteps.shape   ) == 0:
             a__		=			timesteps.astype(dtype=jnp.floataa   )
             a__		=			jnp.expand_dims(__SCREAMING_SNAKE_CASE     ,				0   )

         a__		=			self.time_proj(__SCREAMING_SNAKE_CASE   )
         a__		=			self.time_embedding(__SCREAMING_SNAKE_CASE   )

         # 2. pre-process
         a__		=			jnp.transpose(__SCREAMING_SNAKE_CASE     ,				(0, 2, 3, 1)   )
         a__		=			self.conv_in(__SCREAMING_SNAKE_CASE   )

         a__		=			jnp.transpose(__SCREAMING_SNAKE_CASE     ,				(0, 2, 3, 1)   )
         a__		=			self.controlnet_cond_embedding(__SCREAMING_SNAKE_CASE   )
         sample += controlnet_cond

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

         # 4. mid
         a__		=			self.mid_block(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				deterministic=not train   )

         # 5. contronet blocks
         a__		=			()
         for down_block_res_sample, controlnet_block in zip(__SCREAMING_SNAKE_CASE     ,				self.controlnet_down_blocks   ):
             a__		=			controlnet_block(__SCREAMING_SNAKE_CASE   )
             controlnet_down_block_res_samples += (down_block_res_sample,)

         a__		=			controlnet_down_block_res_samples

         a__		=			self.controlnet_mid_block(__SCREAMING_SNAKE_CASE   )

         # 6. scaling
         a__		=			[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=__SCREAMING_SNAKE_CASE     ,				mid_block_res_sample=__SCREAMING_SNAKE_CASE   )


