

'''simple docstring'''







import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
      from transformers import (
          FlaxBertModel,
          FlaxCLIPVisionModel,
          FlaxVisionTextDualEncoderModel,
          FlaxViTModel,
          VisionTextDualEncoderConfig,
          VisionTextDualEncoderProcessor,
      )
      from transformers.modeling_flax_pytorch_utils import (
          convert_pytorch_state_dict_to_flax,
          load_flax_weights_in_pytorch_model,
      )


if is_torch_available():
      import torch

      from transformers import VisionTextDualEncoderModel

if is_vision_available():
      from PIL import Image
def  UpperCAmelCase__      ( UpperCAmelCase__			)     ->     Dict:
       if isinstance(UpperCAmelCase__,    collections.abc.Iterable			):
              return x
       return (x, x)


@require_flax
class    A__ :



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





           '''simple docstring'''
           pass



    def    snake_case_     (   self    )  -> Tuple:





           '''simple docstring'''
           pass



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





           '''simple docstring'''
           pass



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





           '''simple docstring'''
           A_		   = np.abs((a - b)    ).max()
           self.assertLessEqual(UpperCamelCase__						,			UpperCamelCase__						,			f'''Difference between torch and flax is {diff} (>= {tol}).'''    )



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__=None						,			**UpperCamelCase__    )  -> List[str]:





           '''simple docstring'''
           A_		   = VisionTextDualEncoderConfig.from_vision_text_configs(UpperCamelCase__						,			UpperCamelCase__    )

           A_		   = FlaxVisionTextDualEncoderModel(UpperCamelCase__    )

           A_		   = model(input_ids=UpperCamelCase__						,			pixel_values=UpperCamelCase__						,			attention_mask=UpperCamelCase__    )

           self.assertEqual(output["""text_embeds"""].shape						,			(input_ids.shape[0], config.projection_dim)    )
           self.assertEqual(output["""image_embeds"""].shape						,			(pixel_values.shape[0], config.projection_dim)    )



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__=None						,			**UpperCamelCase__    )  -> str:





           '''simple docstring'''
           A_   ,     A_		   = self.get_vision_text_model(UpperCamelCase__						,			UpperCamelCase__    )
           A_		   = {"""vision_model""": vision_model, """text_model""": text_model}
           A_		   = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCamelCase__    )

           A_		   = model(input_ids=UpperCamelCase__						,			pixel_values=UpperCamelCase__						,			attention_mask=UpperCamelCase__    )

           self.assertEqual(output["""text_embeds"""].shape						,			(input_ids.shape[0], model.config.projection_dim)    )
           self.assertEqual(output["""image_embeds"""].shape						,			(pixel_values.shape[0], model.config.projection_dim)    )



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__=None						,			**UpperCamelCase__    )  -> Optional[Any]:





           '''simple docstring'''
           A_   ,     A_		   = self.get_vision_text_model(UpperCamelCase__						,			UpperCamelCase__    )
           A_		   = {"""vision_model""": vision_model, """text_model""": text_model}
           A_		   = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCamelCase__    )

           A_		   = model(input_ids=UpperCamelCase__						,			pixel_values=UpperCamelCase__						,			attention_mask=UpperCamelCase__    )
           A_		   = output[0]

           with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_pretrained(UpperCamelCase__    )
                  A_		   = FlaxVisionTextDualEncoderModel.from_pretrained(UpperCamelCase__    )

                  A_		   = model(input_ids=UpperCamelCase__						,			pixel_values=UpperCamelCase__						,			attention_mask=UpperCamelCase__    )
                  A_		   = after_output[0]
                  A_		   = np.amax(np.abs(out_a - out_a    )    )
                  self.assertLessEqual(UpperCamelCase__						,			1e-3    )



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__=None						,			**UpperCamelCase__    )  -> int:





           '''simple docstring'''
           A_   ,     A_		   = self.get_vision_text_model(UpperCamelCase__						,			UpperCamelCase__    )
           A_		   = {"""vision_model""": vision_model, """text_model""": text_model}
           A_		   = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCamelCase__    )

           A_		   = model(
               input_ids=UpperCamelCase__						,			pixel_values=UpperCamelCase__						,			attention_mask=UpperCamelCase__						,			output_attentions=UpperCamelCase__    )

           A_		   = output.vision_model_output.attentions
           self.assertEqual(len(UpperCamelCase__    )						,			vision_config.num_hidden_layers    )

           # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
           A_		   = to_atuple(vision_model.config.image_size    )
           A_		   = to_atuple(vision_model.config.patch_size    )
           A_		   = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
           A_		   = num_patches + 1
           self.assertEqual(vision_attentions[0].shape[-3:]						,			(vision_config.num_attention_heads, seq_len, seq_len)    )

           A_		   = output.text_model_output.attentions
           self.assertEqual(len(UpperCamelCase__    )						,			text_config.num_hidden_layers    )

           self.assertEqual(
               text_attentions[0].shape[-3:]						,			(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])						,			)



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





           '''simple docstring'''
           pt_model.to(UpperCamelCase__    )
           pt_model.eval()

           # prepare inputs
           A_		   = inputs_dict
           A_		   = {k: torch.tensor(v.tolist()    ) for k, v in flax_inputs.items()}

           with torch.no_grad():
                  A_		   = pt_model(**UpperCamelCase__    ).to_tuple()

           A_		   = fx_model(**UpperCamelCase__    ).to_tuple()
           self.assertEqual(len(UpperCamelCase__    )						,			len(UpperCamelCase__    )						,			"""Output lengths differ between Flax and PyTorch"""    )
           for fx_output, pt_output in zip(fx_outputs[:4]						,			pt_outputs[:4]    ):
                  self.assert_almost_equals(UpperCamelCase__						,			pt_output.numpy()						,			4e-2    )

           # PT -> Flax
           with tempfile.TemporaryDirectory() as tmpdirname:
                  pt_model.save_pretrained(UpperCamelCase__    )
                  A_		   = FlaxVisionTextDualEncoderModel.from_pretrained(UpperCamelCase__						,			from_pt=UpperCamelCase__    )

           A_		   = fx_model_loaded(**UpperCamelCase__    ).to_tuple()
           self.assertEqual(len(UpperCamelCase__    )						,			len(UpperCamelCase__    )						,			"""Output lengths differ between Flax and PyTorch"""    )
           for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]						,			pt_outputs[:4]    ):
                  self.assert_almost_equals(UpperCamelCase__						,			pt_output.numpy()						,			4e-2    )

           # Flax -> PT
           with tempfile.TemporaryDirectory() as tmpdirname:
                  fx_model.save_pretrained(UpperCamelCase__    )
                  A_		   = VisionTextDualEncoderModel.from_pretrained(UpperCamelCase__						,			from_flax=UpperCamelCase__    )

           pt_model_loaded.to(UpperCamelCase__    )
           pt_model_loaded.eval()

           with torch.no_grad():
                  A_		   = pt_model_loaded(**UpperCamelCase__    ).to_tuple()

           self.assertEqual(len(UpperCamelCase__    )						,			len(UpperCamelCase__    )						,			"""Output lengths differ between Flax and PyTorch"""    )
           for fx_output, pt_output_loaded in zip(fx_outputs[:4]						,			pt_outputs_loaded[:4]    ):
                  self.assert_almost_equals(UpperCamelCase__						,			pt_output_loaded.numpy()						,			4e-2    )



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





           '''simple docstring'''
           A_		   = VisionTextDualEncoderConfig.from_vision_text_configs(UpperCamelCase__						,			UpperCamelCase__    )

           A_		   = VisionTextDualEncoderModel(UpperCamelCase__    )
           A_		   = FlaxVisionTextDualEncoderModel(UpperCamelCase__    )

           A_		   = convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,			UpperCamelCase__    )
           A_		   = fx_state

           self.check_pt_flax_equivalence(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )



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





           '''simple docstring'''
           A_		   = VisionTextDualEncoderConfig.from_vision_text_configs(UpperCamelCase__						,			UpperCamelCase__    )

           A_		   = VisionTextDualEncoderModel(UpperCamelCase__    )
           A_		   = FlaxVisionTextDualEncoderModel(UpperCamelCase__    )

           A_		   = load_flax_weights_in_pytorch_model(UpperCamelCase__						,			fx_model.params    )

           self.check_pt_flax_equivalence(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )



    def    snake_case_     (   self    )  -> Dict:





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



    def    snake_case_     (   self    )  -> Any:





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



    def    snake_case_     (   self    )  -> Dict:





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



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





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



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





           '''simple docstring'''
           A_		   = self.prepare_config_and_inputs()
           A_		   = config_inputs_dict.pop("""vision_config"""    )
           A_		   = config_inputs_dict.pop("""text_config"""    )

           A_		   = config_inputs_dict

           self.check_equivalence_pt_to_flax(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )
           self.check_equivalence_flax_to_pt(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )







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





           '''simple docstring'''
           A_   ,     A_		   = self.get_pretrained_model_and_inputs()

           A_		   = model_a(**UpperCamelCase__    )
           A_		   = outputs[0]

           with tempfile.TemporaryDirectory() as tmp_dirname:
                  model_a.save_pretrained(UpperCamelCase__    )
                  A_		   = FlaxVisionTextDualEncoderModel.from_pretrained(UpperCamelCase__    )

                  A_		   = model_a(**UpperCamelCase__    )
                  A_		   = after_outputs[0]
                  A_		   = np.amax(np.abs(out_a - out_a    )    )
                  self.assertLessEqual(UpperCamelCase__						,			1e-5    )


@require_flax
class    A__ ( _snake_case  ,  unittest.TestCase     ):



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





           '''simple docstring'''
           A_		   = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
               """hf-internal-testing/tiny-random-vit"""						,			"""hf-internal-testing/tiny-bert"""						,			vision_from_pt=UpperCamelCase__						,			text_from_pt=UpperCamelCase__						,			)
           A_		   = 13
           A_		   = floats_tensor(
               [
                   batch_size,
                   model.config.vision_config.num_channels,
                   model.config.vision_config.image_size,
                   model.config.vision_config.image_size,
               ]    )
           A_		   = ids_tensor([batch_size, 4]						,			model.config.text_config.vocab_size    )
           A_		   = random_attention_mask([batch_size, 4]    )
           A_		   = {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

           return model, inputs



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





           '''simple docstring'''
           A_		   = FlaxViTModel(UpperCamelCase__    )
           A_		   = FlaxBertModel(UpperCamelCase__    )
           return vision_model, text_model







    def    snake_case_     (   self    )  -> str:





           '''simple docstring'''
           A_		   = FlaxViTModelTester(self    )
           A_		   = FlaxBertModelTester(self    )
           A_		   = vit_model_tester.prepare_config_and_inputs()
           A_		   = bert_model_tester.prepare_config_and_inputs()

           A_   ,     A_		   = vision_config_and_inputs

           A_   ,     A_   ,     A_   ,     A_		   = text_config_and_inputs

           # make sure that cross attention layers are added
           return {
               "text_config": text_config,
               "vision_config": vision_config,
               "pixel_values": pixel_values,
               "attention_mask": attention_mask,
               "input_ids": input_ids,
               "token_type_ids": token_type_ids,
           }


@require_torch
class    A__ ( _snake_case  ,  unittest.TestCase     ):



    def    snake_case_     (   self    )  -> Any:





           '''simple docstring'''
           A_		   = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
               """hf-internal-testing/tiny-random-clip"""						,			"""hf-internal-testing/tiny-bert"""						,			vision_from_pt=UpperCamelCase__						,			text_from_pt=UpperCamelCase__						,			)
           A_		   = 13
           A_		   = floats_tensor(
               [
                   batch_size,
                   model.config.vision_config.num_channels,
                   model.config.vision_config.image_size,
                   model.config.vision_config.image_size,
               ]    )
           A_		   = ids_tensor([batch_size, 4]						,			model.config.text_config.vocab_size    )
           A_		   = random_attention_mask([batch_size, 4]    )
           A_		   = {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

           return model, inputs



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





           '''simple docstring'''
           A_		   = FlaxCLIPVisionModel(UpperCamelCase__    )
           A_		   = FlaxBertModel(UpperCamelCase__    )
           return vision_model, text_model







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





           '''simple docstring'''
           A_		   = FlaxCLIPVisionModelTester(self    )
           A_		   = FlaxBertModelTester(self    )
           A_		   = clip_model_tester.prepare_config_and_inputs()
           A_		   = bert_model_tester.prepare_config_and_inputs()

           A_   ,     A_		   = vision_config_and_inputs

           A_   ,     A_   ,     A_   ,     A_		   = text_config_and_inputs

           # make sure that cross attention layers are added
           return {
               "text_config": text_config,
               "vision_config": vision_config,
               "pixel_values": pixel_values,
               "attention_mask": attention_mask,
               "input_ids": input_ids,
               "token_type_ids": token_type_ids,
           }


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



    @slow
    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           A_		   = FlaxVisionTextDualEncoderModel.from_pretrained("""clip-italian/clip-italian"""						,			logit_scale_init_value=1.0    )
           A_		   = VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""    )

           A_		   = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
           A_		   = processor(
               text=["""una foto di un gatto""", """una foto di un cane"""]						,			images=UpperCamelCase__						,			padding=UpperCamelCase__						,			return_tensors="""np"""    )

           A_		   = model(**UpperCamelCase__    )

           # verify the logits
           self.assertEqual(outputs.logits_per_image.shape						,			(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])    )
           self.assertEqual(
               outputs.logits_per_text.shape						,			(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])						,			)

           A_		   = np.array([[1.2284727, 0.3104122]]    )

           self.assertTrue(np.allclose(outputs.logits_per_image						,			UpperCamelCase__						,			atol=1e-3    )    )



'''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__    )
