




"""simple docstring"""




from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''
				__UpperCamelCase         =		42
				__UpperCamelCase         =		42






class        __magic_name__				(			nn.Module							):






				'''simple docstring'''
				__UpperCamelCase         =		42
				__UpperCamelCase         =		(16, 32, 96, 2_56)
				__UpperCamelCase         =		jnp.floataa





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   nn.Conv(
							    self.block_out_channels[0]						,					kernel_size=(3, 3)						,					padding=((1, 1), (1, 1))						,					dtype=self.dtype						,					)

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

							lowerCamelCase      =   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						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   self.conv_in(_a					)
							lowerCamelCase      =   nn.silu(_a					)

							for block in self.blocks:
										lowerCamelCase      =   block(_a					)
										lowerCamelCase      =   nn.silu(_a					)

							lowerCamelCase      =   self.conv_out(_a					)

							return embedding






@flax_register_to_config
class        __magic_name__				(			nn.Module		,   UpperCAmelCase__		,   UpperCAmelCase__							):






				'''simple docstring'''
				__UpperCamelCase         =		32
				__UpperCamelCase         =		4
				__UpperCamelCase         =		(
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				__UpperCamelCase         =		False
				__UpperCamelCase         =		(3_20, 6_40, 12_80, 12_80)
				__UpperCamelCase         =		2
				__UpperCamelCase         =		8
				__UpperCamelCase         =		None
				__UpperCamelCase         =		12_80
				__UpperCamelCase         =		0.0
				__UpperCamelCase         =		False
				__UpperCamelCase         =		jnp.floataa
				__UpperCamelCase         =		True
				__UpperCamelCase         =		0
				__UpperCamelCase         =		"rgb"
				__UpperCamelCase         =		(16, 32, 96, 2_56)





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							# init input tensors
							lowerCamelCase      =   (1, self.in_channels, self.sample_size, self.sample_size)
							lowerCamelCase      =   jnp.zeros(_a						,					dtype=jnp.floataa					)
							lowerCamelCase      =   jnp.ones((1,)						,					dtype=jnp.intaa					)
							lowerCamelCase      =   jnp.zeros((1, 1, self.cross_attention_dim)						,					dtype=jnp.floataa					)
							lowerCamelCase      =   (1, 3, self.sample_size * 8, self.sample_size * 8)
							lowerCamelCase      =   jnp.zeros(_a						,					dtype=jnp.floataa					)

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

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





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.block_out_channels
							lowerCamelCase      =   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.
							lowerCamelCase      =   self.num_attention_heads or self.attention_head_dim

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

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

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

							lowerCamelCase      =   self.only_cross_attention
							if isinstance(_a						,					_a					):
										lowerCamelCase      =   (only_cross_attention,) * len(self.down_block_types					)

							if isinstance(_a						,					_a					):
										lowerCamelCase      =   (num_attention_heads,) * len(self.down_block_types					)

							# down
							lowerCamelCase      =   []
							lowerCamelCase      =   []

							lowerCamelCase      =   block_out_channels[0]

							lowerCamelCase      =   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					):
										lowerCamelCase      =   output_channel
										lowerCamelCase      =   block_out_channels[i]
										lowerCamelCase      =   i == len(_a					) - 1

										if down_block_type == "CrossAttnDownBlock2D":
													lowerCamelCase      =   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:
													lowerCamelCase      =   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					):
													lowerCamelCase      =   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:
													lowerCamelCase      =   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					)

							lowerCamelCase      =   down_blocks
							lowerCamelCase      =   controlnet_down_blocks

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

							lowerCamelCase      =   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						,					_a						,					_a						,					_a						,					_a						,					_a = 1.0						,					_a = True						,					_a = False						,					):





							"""simple docstring"""
							lowerCamelCase      =   self.controlnet_conditioning_channel_order
							if channel_order == "bgr":
										lowerCamelCase      =   jnp.flip(_a						,					axis=1					)

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

							lowerCamelCase      =   self.time_proj(_a					)
							lowerCamelCase      =   self.time_embedding(_a					)

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

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

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

							# 4. mid
							lowerCamelCase      =   self.mid_block(_a						,					_a						,					_a						,					deterministic=not train					)

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

							lowerCamelCase      =   controlnet_down_block_res_samples

							lowerCamelCase      =   self.controlnet_mid_block(_a					)

							# 6. scaling
							lowerCamelCase      =   [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					)









"""simple docstring"""




import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowerCAmelCase				:     List[str]					=				object()

# For specifying empty leaf dict `{}`
lowerCAmelCase				:     Any					=				object()



def        a__					(							snake_case__			,					snake_case__							)		->   List[Any]:
			lowerCamelCase      =   tuple((re.compile(x + """$"""							) for x in qs)							)
			for i in range(len(snake_case__							) - len(snake_case__							) + 1							):
						lowerCamelCase      =   [x.match(snake_case__							) for x, y in zip(snake_case__			,					ks[i:]							)]
						if matches and all(snake_case__							):
									return True
			return False



def        a__					(							snake_case__							)		->   str:
			def replace(snake_case__			,					snake_case__							):
						for rule, replacement in rules:
									if _match(snake_case__			,					snake_case__							):
												return replacement
						return val

			return replace



def        a__					(							)		->   Union[str, Any]:
			return [
			    # embeddings
			    (("transformer", "wpe", "embedding"), P("""mp"""			,					snake_case__							)),
			    (("transformer", "wte", "embedding"), P("""mp"""			,					snake_case__							)),
			    # atention
			    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(snake_case__			,					"""mp"""							)),
			    (("attention", "out_proj", "kernel"), P("""mp"""			,					snake_case__							)),
			    (("attention", "out_proj", "bias"), None),
			    # mlp
			    (("mlp", "c_fc", "kernel"), P(snake_case__			,					"""mp"""							)),
			    (("mlp", "c_fc", "bias"), P("""mp"""							)),
			    (("mlp", "c_proj", "kernel"), P("""mp"""			,					snake_case__							)),
			    (("mlp", "c_proj", "bias"), None),
			    # layer norms
			    ((r"ln_\d+", "bias"), None),
			    ((r"\d+", r"ln_\d+", "scale"), None),
			    (("ln_f", "bias"), None),
			    (("ln_f", "scale"), None),
			]



def        a__					(							snake_case__							)		->   Optional[Any]:
			lowerCamelCase      =   _get_partition_rules()
			lowerCamelCase      =   _replacement_rules(snake_case__							)
			lowerCamelCase      =   {k: _unmatched for k in flatten_dict(snake_case__							)}
			lowerCamelCase      =   {k: replace(snake_case__			,					snake_case__							) for k, v in initd.items()}
			assert _unmatched not in result.values(), "Incomplete partition spec."
			return freeze(unflatten_dict(snake_case__							)							)



