
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     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			__UpperCamelCase :  jnp.ndarray
			__UpperCamelCase :  jnp.ndarray




class     _UpperCamelCase			(   nn.Module					):






			'''simple docstring'''

			__UpperCamelCase :  int
			__UpperCamelCase :  Tuple[int]    =	(16, 32, 96, 256)
			__UpperCamelCase :  jnp.dtype    =	jnp.floataa

			def  lowerCAmelCase__		(    self	:    Optional[int] ):
							UpperCamelCase_:							int		=  nn.Conv(
							    self.block_out_channels[0]  ,		kernel_size=(3, 3)  ,		padding=((1, 1), (1, 1))  ,		dtype=self.dtype  ,		)

							UpperCamelCase_:							List[Any]		=  []
							for i in range(len(self.block_out_channels ) - 1 ):
											UpperCamelCase_:							int		=  self.block_out_channels[i]
											UpperCamelCase_:							str		=  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_:							Tuple		=  nn.Conv(
											    snake_case_  ,		kernel_size=(3, 3)  ,		strides=(2, 2)  ,		padding=((1, 1), (1, 1))  ,		dtype=self.dtype  ,		)
											blocks.append(snake_case_ )
							UpperCamelCase_:							List[str]		=  blocks

							UpperCamelCase_:							List[Any]		=  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	:    Optional[Any]  ,		snake_case_	:    List[str] ):
							UpperCamelCase_:							Any		=  self.conv_in(snake_case_ )
							UpperCamelCase_:							List[str]		=  nn.silu(snake_case_ )

							for block in self.blocks:
											UpperCamelCase_:							int		=  block(snake_case_ )
											UpperCamelCase_:							Union[str, Any]		=  nn.silu(snake_case_ )

							UpperCamelCase_:							Any		=  self.conv_out(snake_case_ )

							return embedding




@flax_register_to_config
class     _UpperCamelCase			(   nn.Module       ,    _A       ,    _A					):






			'''simple docstring'''

			__UpperCamelCase :  int    =	32
			__UpperCamelCase :  int    =	4
			__UpperCamelCase :  Tuple[str]    =	(
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "DownBlock2D",
			)
			__UpperCamelCase :  Union[bool, Tuple[bool]]    =	False
			__UpperCamelCase :  Tuple[int]    =	(320, 640, 1_280, 1_280)
			__UpperCamelCase :  int    =	2
			__UpperCamelCase :  Union[int, Tuple[int]]    =	8
			__UpperCamelCase :  Optional[Union[int, Tuple[int]]]    =	None
			__UpperCamelCase :  int    =	1_280
			__UpperCamelCase :  float    =	0.0
			__UpperCamelCase :  bool    =	False
			__UpperCamelCase :  jnp.dtype    =	jnp.floataa
			__UpperCamelCase :  bool    =	True
			__UpperCamelCase :  int    =	0
			__UpperCamelCase :  str    =	"rgb"
			__UpperCamelCase :  Tuple[int]    =	(16, 32, 96, 256)

			def  lowerCAmelCase__		(    self	:    Optional[Any]  ,		snake_case_	:    jax.random.KeyArray ):
							# init input tensors
							UpperCamelCase_:							Optional[int]		=  (1, self.in_channels, self.sample_size, self.sample_size)
							UpperCamelCase_:							List[Any]		=  jnp.zeros(snake_case_  ,		dtype=jnp.floataa )
							UpperCamelCase_:							Dict		=  jnp.ones((1,)  ,		dtype=jnp.intaa )
							UpperCamelCase_:							Any		=  jnp.zeros((1, 1, self.cross_attention_dim)  ,		dtype=jnp.floataa )
							UpperCamelCase_:							Union[str, Any]		=  (1, 3, self.sample_size * 8, self.sample_size * 8)
							UpperCamelCase_:							Union[str, Any]		=  jnp.zeros(snake_case_  ,		dtype=jnp.floataa )

							UpperCamelCase_,				UpperCamelCase_:							str		=  jax.random.split(snake_case_ )
							UpperCamelCase_:							Optional[int]		=  {"""params""": params_rng, """dropout""": dropout_rng}

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

			def  lowerCAmelCase__		(    self	:    Optional[Any] ):
							UpperCamelCase_:							Union[str, Any]		=  self.block_out_channels
							UpperCamelCase_:							int		=  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_:							Union[str, Any]		=  self.num_attention_heads or self.attention_head_dim

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

							# time
							UpperCamelCase_:							List[Any]		=  FlaxTimesteps(
							    block_out_channels[0]  ,		flip_sin_to_cos=self.flip_sin_to_cos  ,		freq_shift=self.config.freq_shift )
							UpperCamelCase_:							Optional[int]		=  FlaxTimestepEmbedding(snake_case_  ,		dtype=self.dtype )

							UpperCamelCase_:							str		=  FlaxControlNetConditioningEmbedding(
							    conditioning_embedding_channels=block_out_channels[0]  ,		block_out_channels=self.conditioning_embedding_out_channels  ,		)

							UpperCamelCase_:							Union[str, Any]		=  self.only_cross_attention
							if isinstance(snake_case_  ,		snake_case_ ):
											UpperCamelCase_:							Optional[Any]		=  (only_cross_attention,) * len(self.down_block_types )

							if isinstance(snake_case_  ,		snake_case_ ):
											UpperCamelCase_:							List[Any]		=  (num_attention_heads,) * len(self.down_block_types )

							# down
							UpperCamelCase_:							Optional[Any]		=  []
							UpperCamelCase_:							Tuple		=  []

							UpperCamelCase_:							List[Any]		=  block_out_channels[0]

							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_ )

							for i, down_block_type in enumerate(self.down_block_types ):
											UpperCamelCase_:							List[str]		=  output_channel
											UpperCamelCase_:							Tuple		=  block_out_channels[i]
											UpperCamelCase_:							Tuple		=  i == len(snake_case_ ) - 1

											if down_block_type == "CrossAttnDownBlock2D":
															UpperCamelCase_:							str		=  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_:							Dict		=  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_:							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_ )

											if not is_final_block:
															UpperCamelCase_:							Dict		=  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_:							Dict		=  down_blocks
							UpperCamelCase_:							Dict		=  controlnet_down_blocks

							# mid
							UpperCamelCase_:							Optional[int]		=  block_out_channels[-1]
							UpperCamelCase_:							str		=  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[Any]		=  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	:    Any  ,		snake_case_	:    Optional[Any]  ,		snake_case_	:    Optional[int]  ,		snake_case_	:    int  ,		snake_case_	:    Optional[Any]  ,		snake_case_	:    float = 1.0  ,		snake_case_	:    bool = True  ,		snake_case_	:    bool = False  ,		):
							UpperCamelCase_:							Optional[Any]		=  self.controlnet_conditioning_channel_order
							if channel_order == "bgr":
											UpperCamelCase_:							List[str]		=  jnp.flip(snake_case_  ,		axis=1 )

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

							UpperCamelCase_:							Optional[Any]		=  self.time_proj(snake_case_ )
							UpperCamelCase_:							Optional[Any]		=  self.time_embedding(snake_case_ )

							# 2. pre-process
							UpperCamelCase_:							Any		=  jnp.transpose(snake_case_  ,		(0, 2, 3, 1) )
							UpperCamelCase_:							Dict		=  self.conv_in(snake_case_ )

							UpperCamelCase_:							int		=  jnp.transpose(snake_case_  ,		(0, 2, 3, 1) )
							UpperCamelCase_:							int		=  self.controlnet_cond_embedding(snake_case_ )
							sample += controlnet_cond

							# 3. down
							UpperCamelCase_:							List[str]		=  (sample,)
							for down_block in self.down_blocks:
											if isinstance(snake_case_  ,		snake_case_ ):
															UpperCamelCase_,				UpperCamelCase_:							Tuple		=  down_block(snake_case_  ,		snake_case_  ,		snake_case_  ,		deterministic=not train )
											else:
															UpperCamelCase_,				UpperCamelCase_:							int		=  down_block(snake_case_  ,		snake_case_  ,		deterministic=not train )
											down_block_res_samples += res_samples

							# 4. mid
							UpperCamelCase_:							Any		=  self.mid_block(snake_case_  ,		snake_case_  ,		snake_case_  ,		deterministic=not train )

							# 5. contronet blocks
							UpperCamelCase_:							Optional[Any]		=  ()
							for down_block_res_sample, controlnet_block in zip(snake_case_  ,		self.controlnet_down_blocks ):
											UpperCamelCase_:							Optional[int]		=  controlnet_block(snake_case_ )
											controlnet_down_block_res_samples += (down_block_res_sample,)

							UpperCamelCase_:							int		=  controlnet_down_block_res_samples

							UpperCamelCase_:							Tuple		=  self.controlnet_mid_block(snake_case_ )

							# 6. scaling
							UpperCamelCase_:							Dict		=  [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_ )








import json
import os
import unittest

from transformers.models.ctrl.tokenization_ctrl import VOCAB_FILES_NAMES, CTRLTokenizer

from ...test_tokenization_common import TokenizerTesterMixin

class     _UpperCamelCase			(   _A       ,    unittest.TestCase					):






			'''simple docstring'''

			__UpperCamelCase :  Tuple    =	CTRLTokenizer
			__UpperCamelCase :  int    =	False
			__UpperCamelCase :  List[str]    =	False

			def  lowerCAmelCase__		(    self	:    List[str] ):
							super().setUp()

							# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
							UpperCamelCase_:							int		=  ["""adapt""", """re@@""", """a@@""", """apt""", """c@@""", """t""", """<unk>"""]
							UpperCamelCase_:							int		=  dict(zip(snake_case_  ,		range(len(snake_case_ ) ) ) )
							UpperCamelCase_:							Union[str, Any]		=  ["""#version: 0.2""", """a p""", """ap t</w>""", """r e""", """a d""", """ad apt</w>""", """"""]
							UpperCamelCase_:							Tuple		=  {"""unk_token""": """<unk>"""}

							UpperCamelCase_:							Optional[Any]		=  os.path.join(self.tmpdirname  ,		VOCAB_FILES_NAMES["""vocab_file"""] )
							UpperCamelCase_:							Dict		=  os.path.join(self.tmpdirname  ,		VOCAB_FILES_NAMES["""merges_file"""] )
							with open(self.vocab_file  ,		"""w"""  ,		encoding="""utf-8""" ) as fp:
											fp.write(json.dumps(snake_case_ ) + """\n""" )
							with open(self.merges_file  ,		"""w"""  ,		encoding="""utf-8""" ) as fp:
											fp.write("""\n""".join(snake_case_ ) )

			def  lowerCAmelCase__		(    self	:    Optional[int]  ,		**snake_case_	:    int ):
							kwargs.update(self.special_tokens_map )
							return CTRLTokenizer.from_pretrained(self.tmpdirname  ,		**snake_case_ )

			def  lowerCAmelCase__		(    self	:    Optional[int]  ,		snake_case_	:    List[str] ):
							UpperCamelCase_:							Dict		=  """adapt react readapt apt"""
							UpperCamelCase_:							List[str]		=  """adapt react readapt apt"""
							return input_text, output_text




			def  lowerCAmelCase__		(    self	:    List[str] ):
							UpperCamelCase_:							str		=  CTRLTokenizer(self.vocab_file  ,		self.merges_file  ,		**self.special_tokens_map )
							UpperCamelCase_:							List[Any]		=  """adapt react readapt apt"""
							UpperCamelCase_:							Optional[int]		=  """adapt re@@ a@@ c@@ t re@@ adapt apt""".split()
							UpperCamelCase_:							int		=  tokenizer.tokenize(snake_case_ )
							self.assertListEqual(snake_case_  ,		snake_case_ )

							UpperCamelCase_:							List[Any]		=  tokens + [tokenizer.unk_token]

							UpperCamelCase_:							Union[str, Any]		=  [0, 1, 2, 4, 5, 1, 0, 3, 6]
							self.assertListEqual(tokenizer.convert_tokens_to_ids(snake_case_ )  ,		snake_case_ )






