

import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, 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 MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
      from transformers.models.mobilevitva.modeling_mobilevitva import (
          MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
          make_divisible,
      )


if is_vision_available():
      from PIL import Image

      from transformers import MobileViTImageProcessor


class    A__ ( _snake_case     ):



    def    snake_case_     (   self    )  -> Any:





           '''simple docstring'''
           A_		   = self.config_class(**self.inputs_dict    )
           self.parent.assertTrue(hasattr(UpperCamelCase__						,			"""width_multiplier"""    )    )


class    A__ :



    def __init__(   self						,			UpperCamelCase__						,			UpperCamelCase__=13						,			UpperCamelCase__=64						,			UpperCamelCase__=2						,			UpperCamelCase__=3						,			UpperCamelCase__="swish"						,			UpperCamelCase__=3						,			UpperCamelCase__=32						,			UpperCamelCase__=0.1						,			UpperCamelCase__=0.02						,			UpperCamelCase__=True						,			UpperCamelCase__=True						,			UpperCamelCase__=10						,			UpperCamelCase__=None						,			UpperCamelCase__=0.25						,			UpperCamelCase__=0.0						,			UpperCamelCase__=0.0						,			)  -> Tuple:





           '''simple docstring'''
           A_		   = parent
           A_		   = batch_size
           A_		   = image_size
           A_		   = patch_size
           A_		   = num_channels
           A_		   = make_divisible(512 * width_multiplier						,			divisor=8    )
           A_		   = hidden_act
           A_		   = conv_kernel_size
           A_		   = output_stride
           A_		   = classifier_dropout_prob
           A_		   = use_labels
           A_		   = is_training
           A_		   = num_labels
           A_		   = initializer_range
           A_		   = scope
           A_		   = width_multiplier
           A_		   = ffn_dropout
           A_		   = attn_dropout



    def    snake_case_     (   self    )  -> int:





           '''simple docstring'''
           A_		   = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

           A_		   = None
           A_		   = None
           if self.use_labels:
                  A_		   = ids_tensor([self.batch_size]						,			self.num_labels    )
                  A_		   = ids_tensor([self.batch_size, self.image_size, self.image_size]						,			self.num_labels    )

           A_		   = self.get_config()

           return config, pixel_values, labels, pixel_labels



    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           return MobileViTVaConfig(
               image_size=self.image_size						,			patch_size=self.patch_size						,			num_channels=self.num_channels						,			hidden_act=self.hidden_act						,			conv_kernel_size=self.conv_kernel_size						,			output_stride=self.output_stride						,			classifier_dropout_prob=self.classifier_dropout_prob						,			initializer_range=self.initializer_range						,			width_multiplier=self.width_multiplier						,			ffn_dropout=self.ffn_dropout_prob						,			attn_dropout=self.attn_dropout_prob						,			)



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





           '''simple docstring'''
           A_		   = MobileViTVaModel(config=UpperCamelCase__    )
           model.to(UpperCamelCase__    )
           model.eval()
           A_		   = model(UpperCamelCase__    )
           self.parent.assertEqual(
               result.last_hidden_state.shape						,			(
                   self.batch_size,
                   self.last_hidden_size,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )						,			)



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )  -> Dict:





           '''simple docstring'''
           A_		   = self.num_labels
           A_		   = MobileViTVaForImageClassification(UpperCamelCase__    )
           model.to(UpperCamelCase__    )
           model.eval()
           A_		   = model(UpperCamelCase__						,			labels=UpperCamelCase__    )
           self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)    )



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





           '''simple docstring'''
           A_		   = self.num_labels
           A_		   = MobileViTVaForSemanticSegmentation(UpperCamelCase__    )
           model.to(UpperCamelCase__    )
           model.eval()
           A_		   = model(UpperCamelCase__    )
           self.parent.assertEqual(
               result.logits.shape						,			(
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )						,			)
           A_		   = model(UpperCamelCase__						,			labels=UpperCamelCase__    )
           self.parent.assertEqual(
               result.logits.shape						,			(
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )						,			)







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





           '''simple docstring'''
           A_		   = self.prepare_config_and_inputs()
           A_   ,     A_   ,     A_   ,     A_		   = config_and_inputs
           A_		   = {"""pixel_values""": pixel_values}
           return config, inputs_dict


@require_torch
class    A__ ( _snake_case  ,  _snake_case  ,  unittest.TestCase     ):
    lowercase		   =						(
        (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
        if is_torch_available()
        else ()
    )

    lowercase		   =						(
        {
            "feature-extraction": MobileViTVaModel,
            "image-classification": MobileViTVaForImageClassification,
            "image-segmentation": MobileViTVaForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    lowercase		   =						False
    lowercase		   =						False
    lowercase		   =						False
    lowercase		   =						False



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





           '''simple docstring'''
           A_		   = MobileViTVaModelTester(self    )
           A_		   = MobileViTVaConfigTester(self						,			config_class=UpperCamelCase__						,			has_text_modality=UpperCamelCase__    )



    def    snake_case_     (   self    )  -> List[Any]:





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



    @unittest.skip(reason="""MobileViTV2 does not use inputs_embeds"""    )
    def    snake_case_     (   self    )  -> Union[str, Any]:





           '''simple docstring'''
           pass



    @unittest.skip(reason="""MobileViTV2 does not support input and output embeddings"""    )
    def    snake_case_     (   self    )  -> Optional[Any]:





           '''simple docstring'''
           pass



    @unittest.skip(reason="""MobileViTV2 does not output attentions"""    )
    def    snake_case_     (   self    )  -> Optional[int]:





           '''simple docstring'''
           pass



    @require_torch_multi_gpu
    @unittest.skip(reason="""Got `CUDA error: misaligned address` for tests after this one being run."""    )
    def    snake_case_     (   self    )  -> Optional[int]:





           '''simple docstring'''
           pass



    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
    def    snake_case_     (   self    )  -> Union[str, Any]:





           '''simple docstring'''
           pass



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





           '''simple docstring'''
           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A_		   = model_class(UpperCamelCase__    )
                  A_		   = inspect.signature(model.forward    )
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  A_		   = [*signature.parameters.keys()]

                  A_		   = ["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]						,			UpperCamelCase__    )



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





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__    )



    def    snake_case_     (   self    )  -> str:





           '''simple docstring'''
           def check_hidden_states_output(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    ):
                  A_		   = model_class(UpperCamelCase__    )
                  model.to(UpperCamelCase__    )
                  model.eval()

                  with torch.no_grad():
                         A_		   = model(**self._prepare_for_class(UpperCamelCase__						,			UpperCamelCase__    )    )

                  A_		   = outputs.hidden_states

                  A_		   = 5
                  self.assertEqual(len(UpperCamelCase__    )						,			UpperCamelCase__    )

                  # MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
                  # with the width and height being successively divided by 2.
                  A_		   = 2
                  for i in range(len(UpperCamelCase__    )    ):
                         self.assertListEqual(
                             list(hidden_states[i].shape[-2:]    )						,			[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]						,			)
                         divisor *= 2

                  self.assertEqual(self.model_tester.output_stride						,			divisor // 2    )

           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A_		   = True
                  check_hidden_states_output(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )

                  # check that output_hidden_states also work using config
                  del inputs_dict["output_hidden_states"]
                  A_		   = True

                  check_hidden_states_output(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )



    def    snake_case_     (   self    )  -> Optional[Any]:





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__    )



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





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase__    )







    @slow
    def    snake_case_     (   self    )  -> List[str]:





           '''simple docstring'''
           for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_		   = MobileViTVaModel.from_pretrained(UpperCamelCase__    )
                  self.assertIsNotNone(UpperCamelCase__    )

def  UpperCAmelCase__      ( )     ->     List[Any]:
       A_		   = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
       return image







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



    @cached_property
    def    snake_case_     (   self    )  -> str:





           '''simple docstring'''
           return (
               MobileViTImageProcessor.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""    )
               if is_vision_available()
               else None
           )



    @slow
    def    snake_case_     (   self    )  -> Union[str, Any]:





           '''simple docstring'''
           A_		   = MobileViTVaForImageClassification.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""    ).to(
               UpperCamelCase__    )

           A_		   = self.default_image_processor
           A_		   = prepare_img()
           A_		   = image_processor(images=UpperCamelCase__						,			return_tensors="""pt"""    ).to(UpperCamelCase__    )

           # forward pass
           with torch.no_grad():
                  A_		   = model(**UpperCamelCase__    )

           # verify the logits
           A_		   = torch.Size((1, 1000)    )
           self.assertEqual(outputs.logits.shape						,			UpperCamelCase__    )

           A_		   = torch.tensor([-1.6_3_3_6e0_0, -7.3_2_0_4e-0_2, -5.1_8_8_3e-0_1]    ).to(UpperCamelCase__    )

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



    @slow
    def    snake_case_     (   self    )  -> Union[str, Any]:





           '''simple docstring'''
           A_		   = MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""    )
           A_		   = model.to(UpperCamelCase__    )

           A_		   = MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""    )

           A_		   = prepare_img()
           A_		   = image_processor(images=UpperCamelCase__						,			return_tensors="""pt"""    ).to(UpperCamelCase__    )

           # forward pass
           with torch.no_grad():
                  A_		   = model(**UpperCamelCase__    )
           A_		   = outputs.logits

           # verify the logits
           A_		   = torch.Size((1, 21, 32, 32)    )
           self.assertEqual(logits.shape						,			UpperCamelCase__    )

           A_		   = torch.tensor(
               [
                   [[7.0863, 7.1525, 6.8201], [6.6931, 6.8770, 6.8933], [6.2978, 7.0366, 6.9636]],
                   [[-3.7134, -3.6712, -3.6675], [-3.5825, -3.3549, -3.4777], [-3.3435, -3.3979, -3.2857]],
                   [[-2.9329, -2.8003, -2.7369], [-3.0564, -2.4780, -2.0207], [-2.6889, -1.9298, -1.7640]],
               ]						,			device=UpperCamelCase__						,			)

           self.assertTrue(torch.allclose(logits[0, :3, :3, :3]						,			UpperCamelCase__						,			atol=1e-4    )    )







    @slow
    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           A_		   = MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""    )
           A_		   = model.to(UpperCamelCase__    )

           A_		   = MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""    )

           A_		   = prepare_img()
           A_		   = image_processor(images=UpperCamelCase__						,			return_tensors="""pt"""    ).to(UpperCamelCase__    )

           # forward pass
           with torch.no_grad():
                  A_		   = model(**UpperCamelCase__    )

           A_		   = outputs.logits.detach().cpu()

           A_		   = image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase__						,			target_sizes=[(50, 60)]    )
           A_		   = torch.Size((50, 60)    )
           self.assertEqual(segmentation[0].shape						,			UpperCamelCase__    )

           A_		   = image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase__    )
           A_		   = torch.Size((32, 32)    )
           self.assertEqual(segmentation[0].shape						,			UpperCamelCase__    )



'''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    A__ ( _snake_case     ):
    lowercase		   =						42
    lowercase		   =						42


class    A__ ( nn.Module     ):
    lowercase		   =						42
    lowercase		   =						(16, 32, 96, 256)
    lowercase		   =						jnp.floataa



    def    snake_case_     (   self    )  -> Any:





           '''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(
                      UpperCamelCase__						,			kernel_size=(3, 3)						,			padding=((1, 1), (1, 1))						,			dtype=self.dtype						,			)
                  blocks.append(UpperCamelCase__    )
                  A_		   = nn.Conv(
                      UpperCamelCase__						,			kernel_size=(3, 3)						,			strides=(2, 2)						,			padding=((1, 1), (1, 1))						,			dtype=self.dtype						,			)
                  blocks.append(UpperCamelCase__    )
           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						,			UpperCamelCase__    )  -> int:





           '''simple docstring'''
           A_		   = self.conv_in(UpperCamelCase__    )
           A_		   = nn.silu(UpperCamelCase__    )

           for block in self.blocks:
                  A_		   = block(UpperCamelCase__    )
                  A_		   = nn.silu(UpperCamelCase__    )

           A_		   = self.conv_out(UpperCamelCase__    )

           return embedding


@flax_register_to_config
class    A__ ( nn.Module  ,  _snake_case  ,  _snake_case     ):
    lowercase		   =						32
    lowercase		   =						4
    lowercase		   =						(
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D",
    )
    lowercase		   =						False
    lowercase		   =						(320, 640, 1_280, 1_280)
    lowercase		   =						2
    lowercase		   =						8
    lowercase		   =						None
    lowercase		   =						1_280
    lowercase		   =						0.0
    lowercase		   =						False
    lowercase		   =						jnp.floataa
    lowercase		   =						True
    lowercase		   =						0
    lowercase		   =						"rgb"
    lowercase		   =						(16, 32, 96, 256)



    def    snake_case_     (   self						,			UpperCamelCase__    )  -> FrozenDict:





           '''simple docstring'''
           # init input tensors
           A_		   = (1, self.in_channels, self.sample_size, self.sample_size)
           A_		   = jnp.zeros(UpperCamelCase__						,			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(UpperCamelCase__						,			dtype=jnp.floataa    )

           A_   ,     A_		   = jax.random.split(UpperCamelCase__    )
           A_		   = {"""params""": params_rng, """dropout""": dropout_rng}

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



    def    snake_case_     (   self    )  -> List[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(UpperCamelCase__						,			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(UpperCamelCase__						,			UpperCamelCase__    ):
                  A_		   = (only_cross_attention,) * len(self.down_block_types    )

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

           # down
           A_		   = []
           A_		   = []

           A_		   = block_out_channels[0]

           A_		   = nn.Conv(
               UpperCamelCase__						,			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(UpperCamelCase__    )

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

                  if down_block_type == "CrossAttnDownBlock2D":
                         A_		   = FlaxCrossAttnDownBlockaD(
                             in_channels=UpperCamelCase__						,			out_channels=UpperCamelCase__						,			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=UpperCamelCase__						,			out_channels=UpperCamelCase__						,			dropout=self.dropout						,			num_layers=self.layers_per_block						,			add_downsample=not is_final_block						,			dtype=self.dtype						,			)

                  down_blocks.append(UpperCamelCase__    )

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

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

           A_		   = down_blocks
           A_		   = controlnet_down_blocks

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

           A_		   = nn.Conv(
               UpperCamelCase__						,			kernel_size=(1, 1)						,			padding="""VALID"""						,			kernel_init=nn.initializers.zeros_init()						,			bias_init=nn.initializers.zeros_init()						,			dtype=self.dtype						,			)







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





           '''simple docstring'''
           A_		   = self.controlnet_conditioning_channel_order
           if channel_order == "bgr":
                  A_		   = jnp.flip(UpperCamelCase__						,			axis=1    )

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

           A_		   = self.time_proj(UpperCamelCase__    )
           A_		   = self.time_embedding(UpperCamelCase__    )

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

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

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

           # 4. mid
           A_		   = self.mid_block(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			deterministic=not train    )

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

           A_		   = controlnet_down_block_res_samples

           A_		   = self.controlnet_mid_block(UpperCamelCase__    )

           # 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=UpperCamelCase__						,			mid_block_res_sample=UpperCamelCase__    )
