

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__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						42
			SCREAMING_SNAKE_CASE							   =						42






class 	__magic_name__    (			nn.Module      ):
			SCREAMING_SNAKE_CASE							   =						42
			SCREAMING_SNAKE_CASE							   =						(1_6, 3_2, 9_6, 2_5_6)
			SCREAMING_SNAKE_CASE							   =						jnp.floataa





			def 		__magic_name__    (		self					)   -> int:




						'''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(
									    __snake_case					,  kernel_size=(3, 3)					,  padding=((1, 1), (1, 1))					,  dtype=self.dtype					,  )
									blocks.append(__snake_case					)
									__a        =nn.Conv(
									    __snake_case					,  kernel_size=(3, 3)					,  strides=(2, 2)					,  padding=((1, 1), (1, 1))					,  dtype=self.dtype					,  )
									blocks.append(__snake_case					)
						__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					,  __snake_case					)   -> Any:




						'''simple docstring'''
						__a        =self.conv_in(__snake_case					)
						__a        =nn.silu(__snake_case					)

						for block in self.blocks:
									__a        =block(__snake_case					)
									__a        =nn.silu(__snake_case					)

						__a        =self.conv_out(__snake_case					)

						return embedding






@flax_register_to_config
class 	__magic_name__    (			nn.Module  ,    lowerCAmelCase_  ,    lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						3_2
			SCREAMING_SNAKE_CASE							   =						4
			SCREAMING_SNAKE_CASE							   =						(
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "DownBlock2D",
			)
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
			SCREAMING_SNAKE_CASE							   =						2
			SCREAMING_SNAKE_CASE							   =						8
			SCREAMING_SNAKE_CASE							   =						None
			SCREAMING_SNAKE_CASE							   =						1_2_8_0
			SCREAMING_SNAKE_CASE							   =						0.0
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						jnp.floataa
			SCREAMING_SNAKE_CASE							   =						True
			SCREAMING_SNAKE_CASE							   =						0
			SCREAMING_SNAKE_CASE							   =						"rgb"
			SCREAMING_SNAKE_CASE							   =						(1_6, 3_2, 9_6, 2_5_6)





			def 		__magic_name__    (		self					,  __snake_case					)   -> FrozenDict:




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

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

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





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




						'''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(__snake_case					,  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(__snake_case					,  __snake_case					):
									__a        =(only_cross_attention,) * len(self.down_block_types					)

						if isinstance(__snake_case					,  __snake_case					):
									__a        =(num_attention_heads,) * len(self.down_block_types					)

						# down
						__a        =[]
						__a        =[]

						__a        =block_out_channels[0]

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

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

						__a        =down_blocks
						__a        =controlnet_down_blocks

						# mid
						__a        =block_out_channels[-1]
						__a        =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					,  )

						__a        =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					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case = 1.0					,  __snake_case = True					,  __snake_case = False					,  )   -> Union[FlaxControlNetOutput, Tuple]:




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

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

						__a        =self.time_proj(__snake_case					)
						__a        =self.time_embedding(__snake_case					)

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

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

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

						# 4. mid
						__a        =self.mid_block(__snake_case					,  __snake_case					,  __snake_case					,  deterministic=not train					)

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

						__a        =controlnet_down_block_res_samples

						__a        =self.controlnet_mid_block(__snake_case					)

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









from manim import *






class 	__magic_name__    (			lowerCAmelCase_      ):





			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						__a        =Rectangle(height=0.5					,  width=0.5					)
						__a        =Rectangle(height=0.25					,  width=0.25					)
						__a        =Rectangle(height=0.46					,  width=0.46					).set_stroke(width=0					)

						__a        =[mem.copy() for i in range(6					)]
						__a        =[mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =Text('CPU'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						cpu.move_to([-2.5, -0.5, 0]					)
						self.add(__snake_case					)

						__a        =[mem.copy() for i in range(4					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =Text('GPU'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						gpu.move_to([-1, -1, 0]					)
						self.add(__snake_case					)

						__a        =[mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)

						__a        =Text('Model'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						model.move_to([3, -1.0, 0]					)
						self.add(__snake_case					)

						__a        =[]
						__a        =[]
						__a        =[]

						for i, rect in enumerate(__snake_case					):
									rect.set_stroke(__snake_case					)

									__a        =Rectangle(height=0.46 / 4					,  width=0.46 / 3					).set_stroke(width=0.0					).set_fill(__snake_case					,  opacity=0.7					)

									if i == 0:
												cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT					)					,  buff=0.02					,  direction=__snake_case					)
												cpu_target.set_x(cpu_target.get_x() + 0.1					)
									elif i == 3:
												cpu_target.next_to(model_cpu_arr[0]					,  direction=__snake_case					,  buff=0.0					)
									else:
												cpu_target.next_to(model_cpu_arr[i - 1]					,  direction=__snake_case					,  buff=0.0					)
									self.add(__snake_case					)
									model_cpu_arr.append(__snake_case					)

						self.add(*__snake_case					,  *__snake_case					,  *__snake_case					)

						__a        =[mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)

						__a        =Text('Loaded Checkpoint'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						checkpoint.move_to([3, 0.5, 0]					)

						self.add(__snake_case					)

						__a        =[]
						__a        =[]

						for i, rect in enumerate(__snake_case					):
									__a        =fill.copy().set_fill(__snake_case					,  opacity=0.7					)
									target.move_to(__snake_case					)
									ckpt_arr.append(__snake_case					)

									__a        =target.copy()
									if i < 5:
												cpu_target.move_to(cpu_left_col_base[i + 1]					)
									else:
												cpu_target.move_to(cpu_right_col_base[i - 5]					)
									ckpt_cpu_arr.append(__snake_case					)
						self.add(*__snake_case					,  *__snake_case					)

						__a        =Square(side_length=2.2					)
						key.move_to([-5, 2, 0]					)

						__a        =MarkupText(
						    f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'					,  font_size=18					,  )

						key_text.move_to([-5, 2.4, 0]					)

						self.add(__snake_case					,  __snake_case					)

						__a        =MarkupText(
						    f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'					,  font_size=18					,  )

						blue_text.next_to(__snake_case					,  DOWN * 2.4					,  aligned_edge=key_text.get_left()					)
						self.add(__snake_case					)

						__a        =MarkupText(
						    f'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.'					,  font_size=24					,  )
						step_a.move_to([2, 2, 0]					)

						__a        =[meta_mem.copy() for i in range(6					)]
						__a        =[meta_mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =Text('Disk'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						disk.move_to([-4.0, -1.25, 0]					)
						self.play(Write(__snake_case					,  run_time=3					)					,  Write(__snake_case					,  run_time=1					)					,  Create(__snake_case					,  run_time=1					)					)

						__a        =[]
						for i, rect in enumerate(__snake_case					):
									__a        =rect.copy()
									target.generate_target()
									target.target.move_to(disk_left_col_base[i]					).scale(0.5					)
									animations.append(MoveToTarget(__snake_case					,  run_time=1.5					)					)
						self.play(*__snake_case					)

						self.play(FadeOut(__snake_case					)					)

						__a        =MarkupText(f'Then, the checkpoint is removed from memory\nthrough garbage collection.'					,  font_size=24					)
						step_a.move_to([2, 2, 0]					)

						self.play(Write(__snake_case					,  run_time=3					)					)

						self.play(
						    FadeOut(__snake_case					,  __snake_case					,  *__snake_case					,  *__snake_case					)					,  )

						self.wait()






