



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     __lowerCAmelCase					(    A		):
			UpperCamelCase           =  4_2
			UpperCamelCase           =  4_2







class     __lowerCAmelCase					(    nn.Module		):
			UpperCamelCase           =  4_2
			UpperCamelCase           =  (1_6, 3_2, 9_6, 2_5_6)
			UpperCamelCase           =  jnp.floataa



			def 					_lowerCamelCase (		self     :						List[Any])     ->					Dict:



						"""simple docstring"""






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

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

						_UpperCAmelCase	 =     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     :						Any				,	A     :						Any)     ->					Any:



						"""simple docstring"""






						_UpperCAmelCase	 =     self.conv_in(A)
						_UpperCAmelCase	 =     nn.silu(A)

						for block in self.blocks:
									_UpperCAmelCase	 =     block(A)
									_UpperCAmelCase	 =     nn.silu(A)

						_UpperCAmelCase	 =     self.conv_out(A)

						return embedding







@flax_register_to_config
class     __lowerCAmelCase					(    nn.Module       ,		A       ,		A		):
			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     :						str				,	A     :						jax.random.KeyArray)     ->					FrozenDict:



						"""simple docstring"""






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

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

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



			def 					_lowerCamelCase (		self     :						str)     ->					Dict:



						"""simple docstring"""






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

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

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

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

						_UpperCAmelCase	 =     self.only_cross_attention
						if isinstance(A				,	A):
									_UpperCAmelCase	 =     (only_cross_attention,) * len(self.down_block_types)

						if isinstance(A				,	A):
									_UpperCAmelCase	 =     (num_attention_heads,) * len(self.down_block_types)

						# down
						_UpperCAmelCase	 =     []
						_UpperCAmelCase	 =     []

						_UpperCAmelCase	 =     block_out_channels[0]

						_UpperCAmelCase	 =     nn.Conv(
						    A				,	kernel_size=(1, 1)				,	padding='VALID'				,	kernel_init=nn.initializers.zeros_init()				,	bias_init=nn.initializers.zeros_init()				,	dtype=self.dtype				,	)
						controlnet_down_blocks.append(A)

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

									if down_block_type == "CrossAttnDownBlock2D":
												_UpperCAmelCase	 =     FlaxCrossAttnDownBlockaD(
												    in_channels=A				,	out_channels=A				,	dropout=self.dropout				,	num_layers=self.layers_per_block				,	num_attention_heads=num_attention_heads[i]				,	add_downsample=not is_final_block				,	use_linear_projection=self.use_linear_projection				,	only_cross_attention=only_cross_attention[i]				,	dtype=self.dtype				,	)
									else:
												_UpperCAmelCase	 =     FlaxDownBlockaD(
												    in_channels=A				,	out_channels=A				,	dropout=self.dropout				,	num_layers=self.layers_per_block				,	add_downsample=not is_final_block				,	dtype=self.dtype				,	)

									down_blocks.append(A)

									for _ in range(self.layers_per_block):
												_UpperCAmelCase	 =     nn.Conv(
												    A				,	kernel_size=(1, 1)				,	padding='VALID'				,	kernel_init=nn.initializers.zeros_init()				,	bias_init=nn.initializers.zeros_init()				,	dtype=self.dtype				,	)
												controlnet_down_blocks.append(A)

									if not is_final_block:
												_UpperCAmelCase	 =     nn.Conv(
												    A				,	kernel_size=(1, 1)				,	padding='VALID'				,	kernel_init=nn.initializers.zeros_init()				,	bias_init=nn.initializers.zeros_init()				,	dtype=self.dtype				,	)
												controlnet_down_blocks.append(A)

						_UpperCAmelCase	 =     down_blocks
						_UpperCAmelCase	 =     controlnet_down_blocks

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

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



			def __call__(		self     :						Optional[Any]				,	A     :						Any				,	A     :						Any				,	A     :						Any				,	A     :						List[str]				,	A     :						float = 1.0				,	A     :						bool = True				,	A     :						bool = False				,	)     ->					Union[FlaxControlNetOutput, Tuple]:



						"""simple docstring"""






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

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

						_UpperCAmelCase	 =     self.time_proj(A)
						_UpperCAmelCase	 =     self.time_embedding(A)

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

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

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

						# 4. mid
						_UpperCAmelCase	 =     self.mid_block(A				,	A				,	A				,	deterministic=not train)

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

						_UpperCAmelCase	 =     controlnet_down_block_res_samples

						_UpperCAmelCase	 =     self.controlnet_mid_block(A)

						# 6. scaling
						_UpperCAmelCase	 =     [sample * conditioning_scale for sample in down_block_res_samples]
						mid_block_res_sample *= conditioning_scale

						if not return_dict:
									return (down_block_res_samples, mid_block_res_sample)

						return FlaxControlNetOutput(
						    down_block_res_samples=A				,	mid_block_res_sample=A)





import unittest

from transformers import is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
							from PIL import Image
else:







							class     __lowerCAmelCase					:



										@staticmethod
										def 					_lowerCamelCase (		*A     :						Union[str, Any]				,	**A     :						List[Any])     ->					Union[str, Any]:



													"""simple docstring"""






													pass







@is_pipeline_test
@require_vision
class     __lowerCAmelCase					(    unittest.TestCase		):



			@require_torch
			def 					_lowerCamelCase (		self     :						List[str])     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    model='hf-internal-testing/tiny-random-clip-zero-shot-image-classification'				,	)
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['a', 'b', 'c'])

						# The floating scores are so close, we enter floating error approximation and the order is not guaranteed across
						# python and torch versions.
						self.assertIn(
						    nested_simplify(A)				,	[
						        [{'score': 0.3_3_3, 'label': 'a'}, {'score': 0.3_3_3, 'label': 'b'}, {'score': 0.3_3_3, 'label': 'c'}],
						        [{'score': 0.3_3_3, 'label': 'a'}, {'score': 0.3_3_3, 'label': 'c'}, {'score': 0.3_3_3, 'label': 'b'}],
						    ]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['A', 'B', 'C']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						    ]				,	)



			@require_tf
			def 					_lowerCamelCase (		self     :						str)     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    model='hf-internal-testing/tiny-random-clip-zero-shot-image-classification'				,	framework='tf')
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['a', 'b', 'c'])

						self.assertEqual(
						    nested_simplify(A)				,	[{'score': 0.3_3_3, 'label': 'a'}, {'score': 0.3_3_3, 'label': 'b'}, {'score': 0.3_3_3, 'label': 'c'}]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['A', 'B', 'C']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						    ]				,	)



			@slow
			@require_torch
			def 					_lowerCamelCase (		self     :						Tuple)     ->					Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    task='zero-shot-image-classification'				,	model='openai/clip-vit-base-patch32'				,	)
						# This is an image of 2 cats with remotes and no planes
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['cat', 'plane', 'remote'])

						self.assertEqual(
						    nested_simplify(A)				,	[
						        {'score': 0.5_1_1, 'label': 'remote'},
						        {'score': 0.4_8_5, 'label': 'cat'},
						        {'score': 0.0_0_4, 'label': 'plane'},
						    ]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['cat', 'plane', 'remote']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.5_1_1, 'label': 'remote'},
						            {'score': 0.4_8_5, 'label': 'cat'},
						            {'score': 0.0_0_4, 'label': 'plane'},
						        ],
						    ]
						    * 5				,	)



			@slow
			@require_tf
			def 					_lowerCamelCase (		self     :						List[str])     ->					Any:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    task='zero-shot-image-classification'				,	model='openai/clip-vit-base-patch32'				,	framework='tf')
						# This is an image of 2 cats with remotes and no planes
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['cat', 'plane', 'remote'])
						self.assertEqual(
						    nested_simplify(A)				,	[
						        {'score': 0.5_1_1, 'label': 'remote'},
						        {'score': 0.4_8_5, 'label': 'cat'},
						        {'score': 0.0_0_4, 'label': 'plane'},
						    ]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['cat', 'plane', 'remote']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.5_1_1, 'label': 'remote'},
						            {'score': 0.4_8_5, 'label': 'cat'},
						            {'score': 0.0_0_4, 'label': 'plane'},
						        ],
						    ]
						    * 5				,	)
