


'''simple docstring'''






import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
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 transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
 from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor






class  __magic_name__							:






       def __init__(      self  ,							snake_case  ,							snake_case=1_3  ,							snake_case=3_2  ,							snake_case=3  ,							snake_case=4  ,							snake_case=[1_0, 2_0, 3_0, 4_0]  ,							snake_case=[2, 2, 3, 2]  ,							snake_case=True  ,							snake_case=True  ,							snake_case=3_7  ,							snake_case="gelu"  ,							snake_case=1_0  ,							snake_case=0.02  ,							snake_case=["stage2", "stage3", "stage4"]  ,							snake_case=[2, 3, 4]  ,							snake_case=None  ,							)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =parent
          _UpperCAmelCase       :       Optional[Any]							   =batch_size
          _UpperCAmelCase       :       Any							   =image_size
          _UpperCAmelCase       :       Any							   =num_channels
          _UpperCAmelCase       :       str							   =num_stages
          _UpperCAmelCase       :       Tuple							   =hidden_sizes
          _UpperCAmelCase       :       Union[str, Any]							   =depths
          _UpperCAmelCase       :       str							   =is_training
          _UpperCAmelCase       :       List[str]							   =use_labels
          _UpperCAmelCase       :       str							   =intermediate_size
          _UpperCAmelCase       :       Union[str, Any]							   =hidden_act
          _UpperCAmelCase       :       Tuple							   =num_labels
          _UpperCAmelCase       :       Tuple							   =initializer_range
          _UpperCAmelCase       :       Union[str, Any]							   =out_features
          _UpperCAmelCase       :       Dict							   =out_indices
          _UpperCAmelCase       :       Optional[Any]							   =scope






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

          _UpperCAmelCase       :       int							   =self.get_config()

          return config, pixel_values, labels






       def       lowerCAmelCase		(      self)						->  Dict:




          '''simple docstring'''
          return ConvNextVaConfig(
              num_channels=self.num_channels  ,							hidden_sizes=self.hidden_sizes  ,							depths=self.depths  ,							num_stages=self.num_stages  ,							hidden_act=self.hidden_act  ,							is_decoder=snake_case  ,							initializer_range=self.initializer_range  ,							out_features=self.out_features  ,							out_indices=self.out_indices  ,							num_labels=self.num_labels  ,							)






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case  ,							snake_case)						->  Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =ConvNextVaModel(config=snake_case)
          model.to(snake_case)
          model.eval()
          _UpperCAmelCase       :       Dict							   =model(snake_case)
          # expected last hidden states: B, C, H // 32, W // 32
          self.parent.assertEqual(
              result.last_hidden_state.shape  ,							(self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2)  ,							)






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case  ,							snake_case)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       List[Any]							   =ConvNextVaForImageClassification(snake_case)
          model.to(snake_case)
          model.eval()
          _UpperCAmelCase       :       List[Any]							   =model(snake_case  ,							labels=snake_case)
          self.parent.assertEqual(result.logits.shape  ,							(self.batch_size, self.num_labels))






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case  ,							snake_case)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =ConvNextVaBackbone(config=snake_case)
          model.to(snake_case)
          model.eval()
          _UpperCAmelCase       :       Optional[int]							   =model(snake_case)

          # verify hidden states
          self.parent.assertEqual(len(result.feature_maps)  ,							len(config.out_features))
          self.parent.assertListEqual(list(result.feature_maps[0].shape)  ,							[self.batch_size, self.hidden_sizes[1], 4, 4])

          # verify channels
          self.parent.assertEqual(len(model.channels)  ,							len(config.out_features))
          self.parent.assertListEqual(model.channels  ,							config.hidden_sizes[1:])

          # verify backbone works with out_features=None
          _UpperCAmelCase       :       Union[str, Any]							   =None
          _UpperCAmelCase       :       Dict							   =ConvNextVaBackbone(config=snake_case)
          model.to(snake_case)
          model.eval()
          _UpperCAmelCase       :       List[str]							   =model(snake_case)

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps)  ,							1)
          self.parent.assertListEqual(list(result.feature_maps[0].shape)  ,							[self.batch_size, self.hidden_sizes[-1], 1, 1])

          # verify channels
          self.parent.assertEqual(len(model.channels)  ,							1)
          self.parent.assertListEqual(model.channels  ,							[config.hidden_sizes[-1]])






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       List[str]							   =self.prepare_config_and_inputs()
          _UpperCAmelCase   ,	_UpperCAmelCase   ,	_UpperCAmelCase       :       Optional[int]							   =config_and_inputs
          _UpperCAmelCase       :       Union[str, Any]							   ={'pixel_values': pixel_values}
          return config, inputs_dict






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =self.prepare_config_and_inputs()
          _UpperCAmelCase   ,	_UpperCAmelCase   ,	_UpperCAmelCase       :       Optional[Any]							   =config_and_inputs
          _UpperCAmelCase       :       Union[str, Any]							   ={'pixel_values': pixel_values, 'labels': labels}
          return config, inputs_dict






@require_torch
class  __magic_name__							(		lowerCAmelCase ,lowerCAmelCase ,unittest.TestCase    ):
       UpperCAmelCase										=(
           (
               ConvNextVaModel,
               ConvNextVaForImageClassification,
               ConvNextVaBackbone,
           )
           if is_torch_available()
           else ()
       )
       UpperCAmelCase										=(
           {"feature-extraction": ConvNextVaModel, "image-classification": ConvNextVaForImageClassification}
           if is_torch_available()
           else {}
       )

       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=False






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




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =ConvNextVaModelTester(self)
          _UpperCAmelCase       :       Any							   =ConfigTester(self  ,							config_class=snake_case  ,							has_text_modality=snake_case  ,							hidden_size=3_7)






       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()






       def       lowerCAmelCase		(      self)						->  int:




          '''simple docstring'''
          return






       @unittest.skip(reason='ConvNextV2 does not use inputs_embeds')
       def       lowerCAmelCase		(      self)						->  List[Any]:




          '''simple docstring'''
          pass






       @unittest.skip(reason='ConvNextV2 does not support input and output embeddings')
       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          pass






       @unittest.skip(reason='ConvNextV2 does not use feedforward chunking')
       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          pass






       def       lowerCAmelCase		(      self)						->  Dict:




          '''simple docstring'''
          if not self.model_tester.is_training:
             return

          for model_class in self.all_model_classes:
             _UpperCAmelCase   ,	_UpperCAmelCase       :       str							   =self.model_tester.prepare_config_and_inputs_with_labels()
             _UpperCAmelCase       :       Dict							   =True

             if model_class.__name__ in [
                 *get_values(snake_case),
                 *get_values(snake_case),
             ]:
                continue

             _UpperCAmelCase       :       Union[str, Any]							   =model_class(snake_case)
             model.to(snake_case)
             model.train()
             _UpperCAmelCase       :       str							   =self._prepare_for_class(snake_case  ,							snake_case  ,							return_labels=snake_case)
             _UpperCAmelCase       :       Tuple							   =model(**snake_case).loss
             loss.backward()






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          if not self.model_tester.is_training:
             return

          for model_class in self.all_model_classes:
             _UpperCAmelCase   ,	_UpperCAmelCase       :       Optional[int]							   =self.model_tester.prepare_config_and_inputs_with_labels()
             _UpperCAmelCase       :       Dict							   =False
             _UpperCAmelCase       :       str							   =True

             if (
                 model_class.__name__
                 in [*get_values(snake_case), *get_values(snake_case)]
                 or not model_class.supports_gradient_checkpointing
             ):
                continue

             _UpperCAmelCase       :       Any							   =model_class(snake_case)
             model.to(snake_case)
             model.gradient_checkpointing_enable()
             model.train()
             _UpperCAmelCase       :       str							   =self._prepare_for_class(snake_case  ,							snake_case  ,							return_labels=snake_case)
             _UpperCAmelCase       :       int							   =model(**snake_case).loss
             loss.backward()






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase   ,	_UpperCAmelCase       :       Dict							   =self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _UpperCAmelCase       :       Optional[Any]							   =model_class(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       :       List[str]							   =['pixel_values']
             self.assertListEqual(arg_names[:1]  ,							snake_case)






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case)






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          def check_hidden_states_output(snake_case  ,							snake_case  ,							snake_case):
             _UpperCAmelCase       :       Dict							   =model_class(snake_case)
             model.to(snake_case)
             model.eval()

             with torch.no_grad():
                _UpperCAmelCase       :       Any							   =model(**self._prepare_for_class(snake_case  ,							snake_case))

             _UpperCAmelCase       :       Any							   =outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

             _UpperCAmelCase       :       Any							   =self.model_tester.num_stages
             self.assertEqual(len(snake_case)  ,							expected_num_stages + 1)

             # ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
             self.assertListEqual(
                 list(hidden_states[0].shape[-2:])  ,							[self.model_tester.image_size // 4, self.model_tester.image_size // 4]  ,							)

          _UpperCAmelCase   ,	_UpperCAmelCase       :       Optional[int]							   =self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _UpperCAmelCase       :       List[Any]							   =True
             check_hidden_states_output(snake_case  ,							snake_case  ,							snake_case)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             _UpperCAmelCase       :       Union[str, Any]							   =True

             check_hidden_states_output(snake_case  ,							snake_case  ,							snake_case)






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




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*snake_case)






       @slow
       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase       :       List[str]							   =ConvNextVaModel.from_pretrained(snake_case)
             self.assertIsNotNone(snake_case)






def lowerCamelCase__      (						):







   '''simple docstring'''




   _UpperCAmelCase       :       Tuple							   =Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'							)
   return image






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






       @cached_property
       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          return AutoImageProcessor.from_pretrained('facebook/convnextv2-tiny-1k-224') if is_vision_available() else None






       @slow
       def       lowerCAmelCase		(      self)						->  Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =ConvNextVaForImageClassification.from_pretrained('facebook/convnextv2-tiny-1k-224').to(snake_case)

          _UpperCAmelCase       :       Dict							   =self.default_image_processor
          _UpperCAmelCase       :       int							   =prepare_img()
          _UpperCAmelCase       :       List[Any]							   =preprocessor(images=snake_case  ,							return_tensors='pt').to(snake_case)

          # forward pass
          with torch.no_grad():
             _UpperCAmelCase       :       int							   =model(**snake_case)

          # verify the logits
          _UpperCAmelCase       :       Optional[int]							   =torch.Size((1, 1_0_0_0))
          self.assertEqual(outputs.logits.shape  ,							snake_case)

          _UpperCAmelCase       :       Optional[int]							   =torch.tensor([0.99_96, 0.19_66, -0.43_86]).to(snake_case)
          self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,							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  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										=42
       UpperCAmelCase										=42






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=42
       UpperCAmelCase										=(1_6, 3_2, 9_6, 2_5_6)
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =nn.Conv(
              self.block_out_channels[0]  ,							kernel_size=(3, 3)  ,							padding=((1, 1), (1, 1))  ,							dtype=self.dtype  ,							)

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

          _UpperCAmelCase       :       Tuple							   =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  ,							snake_case)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =self.conv_in(snake_case)
          _UpperCAmelCase       :       Any							   =nn.silu(snake_case)

          for block in self.blocks:
             _UpperCAmelCase       :       Optional[Any]							   =block(snake_case)
             _UpperCAmelCase       :       Union[str, Any]							   =nn.silu(snake_case)

          _UpperCAmelCase       :       str							   =self.conv_out(snake_case)

          return embedding






@flax_register_to_config
class  __magic_name__							(		nn.Module ,lowerCAmelCase ,lowerCAmelCase    ):
       UpperCAmelCase										=3_2
       UpperCAmelCase										=4
       UpperCAmelCase										=(
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       UpperCAmelCase										=False
       UpperCAmelCase										=(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
       UpperCAmelCase										=2
       UpperCAmelCase										=8
       UpperCAmelCase										=None
       UpperCAmelCase										=1_2_8_0
       UpperCAmelCase										=0.0
       UpperCAmelCase										=False
       UpperCAmelCase										=jnp.floataa
       UpperCAmelCase										=True
       UpperCAmelCase										=0
       UpperCAmelCase										="rgb"
       UpperCAmelCase										=(1_6, 3_2, 9_6, 2_5_6)






       def       lowerCAmelCase		(      self  ,							snake_case)						->  FrozenDict:




          '''simple docstring'''
          # init input tensors
          _UpperCAmelCase       :       Any							   =(1, self.in_channels, self.sample_size, self.sample_size)
          _UpperCAmelCase       :       Optional[Any]							   =jnp.zeros(snake_case  ,							dtype=jnp.floataa)
          _UpperCAmelCase       :       Optional[int]							   =jnp.ones((1,)  ,							dtype=jnp.intaa)
          _UpperCAmelCase       :       str							   =jnp.zeros((1, 1, self.cross_attention_dim)  ,							dtype=jnp.floataa)
          _UpperCAmelCase       :       Optional[Any]							   =(1, 3, self.sample_size * 8, self.sample_size * 8)
          _UpperCAmelCase       :       int							   =jnp.zeros(snake_case  ,							dtype=jnp.floataa)

          _UpperCAmelCase   ,	_UpperCAmelCase       :       List[Any]							   =jax.random.split(snake_case)
          _UpperCAmelCase       :       str							   ={'params': params_rng, 'dropout': dropout_rng}

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






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =self.block_out_channels
          _UpperCAmelCase       :       Tuple							   =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.
          _UpperCAmelCase       :       Optional[Any]							   =self.num_attention_heads or self.attention_head_dim

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

          # time
          _UpperCAmelCase       :       Union[str, Any]							   =FlaxTimesteps(
              block_out_channels[0]  ,							flip_sin_to_cos=self.flip_sin_to_cos  ,							freq_shift=self.config.freq_shift)
          _UpperCAmelCase       :       str							   =FlaxTimestepEmbedding(snake_case  ,							dtype=self.dtype)

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

          _UpperCAmelCase       :       Optional[int]							   =self.only_cross_attention
          if isinstance(snake_case  ,							snake_case):
             _UpperCAmelCase       :       Dict							   =(only_cross_attention,) * len(self.down_block_types)

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

          # down
          _UpperCAmelCase       :       int							   =[]
          _UpperCAmelCase       :       Optional[int]							   =[]

          _UpperCAmelCase       :       List[str]							   =block_out_channels[0]

          _UpperCAmelCase       :       int							   =nn.Conv(
              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(snake_case)

          for i, down_block_type in enumerate(self.down_block_types):
             _UpperCAmelCase       :       Tuple							   =output_channel
             _UpperCAmelCase       :       Dict							   =block_out_channels[i]
             _UpperCAmelCase       :       str							   =i == len(snake_case) - 1

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

             down_blocks.append(snake_case)

             for _ in range(self.layers_per_block):
                _UpperCAmelCase       :       Tuple							   =nn.Conv(
                    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(snake_case)

             if not is_final_block:
                _UpperCAmelCase       :       List[str]							   =nn.Conv(
                    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(snake_case)

          _UpperCAmelCase       :       List[Any]							   =down_blocks
          _UpperCAmelCase       :       Optional[Any]							   =controlnet_down_blocks

          # mid
          _UpperCAmelCase       :       int							   =block_out_channels[-1]
          _UpperCAmelCase       :       Optional[Any]							   =FlaxUNetMidBlockaDCrossAttn(
              in_channels=snake_case  ,							dropout=self.dropout  ,							num_attention_heads=num_attention_heads[-1]  ,							use_linear_projection=self.use_linear_projection  ,							dtype=self.dtype  ,							)

          _UpperCAmelCase       :       Optional[int]							   =nn.Conv(
              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  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = 1.0  ,							snake_case = True  ,							snake_case = False  ,							)						->  Union[FlaxControlNetOutput, Tuple]:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =self.controlnet_conditioning_channel_order
          if channel_order == "bgr":
             _UpperCAmelCase       :       Optional[int]							   =jnp.flip(snake_case  ,							axis=1)

          # 1. time
          if not isinstance(snake_case  ,							jnp.ndarray):
             _UpperCAmelCase       :       Optional[int]							   =jnp.array([timesteps]  ,							dtype=jnp.intaa)
          elif isinstance(snake_case  ,							jnp.ndarray) and len(timesteps.shape) == 0:
             _UpperCAmelCase       :       str							   =timesteps.astype(dtype=jnp.floataa)
             _UpperCAmelCase       :       Dict							   =jnp.expand_dims(snake_case  ,							0)

          _UpperCAmelCase       :       int							   =self.time_proj(snake_case)
          _UpperCAmelCase       :       Any							   =self.time_embedding(snake_case)

          # 2. pre-process
          _UpperCAmelCase       :       str							   =jnp.transpose(snake_case  ,							(0, 2, 3, 1))
          _UpperCAmelCase       :       Any							   =self.conv_in(snake_case)

          _UpperCAmelCase       :       List[str]							   =jnp.transpose(snake_case  ,							(0, 2, 3, 1))
          _UpperCAmelCase       :       Optional[int]							   =self.controlnet_cond_embedding(snake_case)
          sample += controlnet_cond

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

          # 4. mid
          _UpperCAmelCase       :       List[Any]							   =self.mid_block(snake_case  ,							snake_case  ,							snake_case  ,							deterministic=not train)

          # 5. contronet blocks
          _UpperCAmelCase       :       Union[str, Any]							   =()
          for down_block_res_sample, controlnet_block in zip(snake_case  ,							self.controlnet_down_blocks):
             _UpperCAmelCase       :       List[str]							   =controlnet_block(snake_case)
             controlnet_down_block_res_samples += (down_block_res_sample,)

          _UpperCAmelCase       :       Optional[int]							   =controlnet_down_block_res_samples

          _UpperCAmelCase       :       List[str]							   =self.controlnet_mid_block(snake_case)

          # 6. scaling
          _UpperCAmelCase       :       Tuple							   =[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=snake_case  ,							mid_block_res_sample=snake_case)

