




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,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)
@flax.struct.dataclass
class 							__snake_case	(_UpperCamelCase		):
				lowerCAmelCase__      						=  42






@flax_register_to_config
class 							__snake_case	(nn.Module    ,						_UpperCamelCase    ,						_UpperCamelCase		):
				lowerCAmelCase__      						=  3_2
				lowerCAmelCase__      						=  4
				lowerCAmelCase__      						=  4
				lowerCAmelCase__      						=  (
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				lowerCAmelCase__      						=  ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				lowerCAmelCase__      						=  False
				lowerCAmelCase__      						=  (3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
				lowerCAmelCase__      						=  2
				lowerCAmelCase__      						=  8
				lowerCAmelCase__      						=  None
				lowerCAmelCase__      						=  1_2_8_0
				lowerCAmelCase__      						=  0.0
				lowerCAmelCase__      						=  False
				lowerCAmelCase__      						=  jnp.floataa
				lowerCAmelCase__      						=  True
				lowerCAmelCase__      						=  0
				lowerCAmelCase__      						=  False



				def    SCREAMING_SNAKE_CASE							(   self   :		Optional[Any]		,						_UpperCAmelCase   :		str   )	->     List[str]:





								'''simple docstring'''




								_lowerCAmelCase     :     Optional[Any]              =   (1, self.in_channels, self.sample_size, self.sample_size)
								_lowerCAmelCase     :     Optional[int]              =   jnp.zeros(lowercase__		,						dtype=jnp.floataa   )
								_lowerCAmelCase     :     List[str]              =   jnp.ones((1,)		,						dtype=jnp.intaa   )
								_lowerCAmelCase     :     Union[str, Any]              =   jnp.zeros((1, 1, self.cross_attention_dim)		,						dtype=jnp.floataa   )

								_lowerCAmelCase     :     List[Any]              =   jax.random.split(lowercase__   )
								_lowerCAmelCase     :     Union[str, Any]              =   {'''params''': params_rng, '''dropout''': dropout_rng}

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



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





								'''simple docstring'''




								_lowerCAmelCase     :     Any              =   self.block_out_channels
								_lowerCAmelCase     :     Union[str, Any]              =   block_out_channels[0] * 4

								if self.num_attention_heads is not None:
												raise ValueError(
												    """At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."""   )

								# 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     :     Dict              =   self.num_attention_heads or self.attention_head_dim

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

								# time
								_lowerCAmelCase     :     str              =   FlaxTimesteps(
								    block_out_channels[0]		,						flip_sin_to_cos=self.flip_sin_to_cos		,						freq_shift=self.config.freq_shift   )
								_lowerCAmelCase     :     str              =   FlaxTimestepEmbedding(lowercase__		,						dtype=self.dtype   )

								_lowerCAmelCase     :     int              =   self.only_cross_attention
								if isinstance(lowercase__		,						lowercase__   ):
												_lowerCAmelCase     :     List[str]              =   (only_cross_attention,) * len(self.down_block_types   )

								if isinstance(lowercase__		,						lowercase__   ):
												_lowerCAmelCase     :     Tuple              =   (num_attention_heads,) * len(self.down_block_types   )

								# down
								_lowerCAmelCase     :     Any              =   []
								_lowerCAmelCase     :     Optional[Any]              =   block_out_channels[0]
								for i, down_block_type in enumerate(self.down_block_types   ):
												_lowerCAmelCase     :     Optional[Any]              =   output_channel
												_lowerCAmelCase     :     Optional[int]              =   block_out_channels[i]
												_lowerCAmelCase     :     List[str]              =   i == len(lowercase__   ) - 1

												if down_block_type == "CrossAttnDownBlock2D":
																_lowerCAmelCase     :     Optional[Any]              =   FlaxCrossAttnDownBlockaD(
																    in_channels=lowercase__		,						out_channels=lowercase__		,						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]		,						use_memory_efficient_attention=self.use_memory_efficient_attention		,						dtype=self.dtype		,						)
												else:
																_lowerCAmelCase     :     int              =   FlaxDownBlockaD(
																    in_channels=lowercase__		,						out_channels=lowercase__		,						dropout=self.dropout		,						num_layers=self.layers_per_block		,						add_downsample=not is_final_block		,						dtype=self.dtype		,						)

												down_blocks.append(lowercase__   )
								_lowerCAmelCase     :     Optional[Any]              =   down_blocks

								# mid
								_lowerCAmelCase     :     Optional[Any]              =   FlaxUNetMidBlockaDCrossAttn(
								    in_channels=block_out_channels[-1]		,						dropout=self.dropout		,						num_attention_heads=num_attention_heads[-1]		,						use_linear_projection=self.use_linear_projection		,						use_memory_efficient_attention=self.use_memory_efficient_attention		,						dtype=self.dtype		,						)

								# up
								_lowerCAmelCase     :     int              =   []
								_lowerCAmelCase     :     Optional[int]              =   list(reversed(lowercase__   )   )
								_lowerCAmelCase     :     Any              =   list(reversed(lowercase__   )   )
								_lowerCAmelCase     :     Dict              =   list(reversed(lowercase__   )   )
								_lowerCAmelCase     :     Tuple              =   reversed_block_out_channels[0]
								for i, up_block_type in enumerate(self.up_block_types   ):
												_lowerCAmelCase     :     Optional[int]              =   output_channel
												_lowerCAmelCase     :     Any              =   reversed_block_out_channels[i]
												_lowerCAmelCase     :     Optional[Any]              =   reversed_block_out_channels[min(i + 1		,						len(lowercase__   ) - 1   )]

												_lowerCAmelCase     :     int              =   i == len(lowercase__   ) - 1

												if up_block_type == "CrossAttnUpBlock2D":
																_lowerCAmelCase     :     List[Any]              =   FlaxCrossAttnUpBlockaD(
																    in_channels=lowercase__		,						out_channels=lowercase__		,						prev_output_channel=lowercase__		,						num_layers=self.layers_per_block + 1		,						num_attention_heads=reversed_num_attention_heads[i]		,						add_upsample=not is_final_block		,						dropout=self.dropout		,						use_linear_projection=self.use_linear_projection		,						only_cross_attention=only_cross_attention[i]		,						use_memory_efficient_attention=self.use_memory_efficient_attention		,						dtype=self.dtype		,						)
												else:
																_lowerCAmelCase     :     Optional[Any]              =   FlaxUpBlockaD(
																    in_channels=lowercase__		,						out_channels=lowercase__		,						prev_output_channel=lowercase__		,						num_layers=self.layers_per_block + 1		,						add_upsample=not is_final_block		,						dropout=self.dropout		,						dtype=self.dtype		,						)

												up_blocks.append(lowercase__   )
												_lowerCAmelCase     :     int              =   output_channel
								_lowerCAmelCase     :     Optional[Any]              =   up_blocks

								# out
								_lowerCAmelCase     :     Dict              =   nn.GroupNorm(num_groups=32		,						epsilon=1E-5   )
								_lowerCAmelCase     :     List[str]              =   nn.Conv(
								    self.out_channels		,						kernel_size=(3, 3)		,						strides=(1, 1)		,						padding=((1, 1), (1, 1))		,						dtype=self.dtype		,						)



				def __call__(   self   :		Optional[Any]		,						_UpperCAmelCase   :		List[Any]		,						_UpperCAmelCase   :		Optional[Any]		,						_UpperCAmelCase   :		Optional[int]		,						_UpperCAmelCase   :		List[Any]=None		,						_UpperCAmelCase   :		Any=None		,						_UpperCAmelCase   :		str = True		,						_UpperCAmelCase   :		Tuple = False		,						)	->     Dict:





								'''simple docstring'''




								if not isinstance(lowercase__		,						jnp.ndarray   ):
												_lowerCAmelCase     :     Union[str, Any]              =   jnp.array([timesteps]		,						dtype=jnp.intaa   )
								elif isinstance(lowercase__		,						jnp.ndarray   ) and len(timesteps.shape   ) == 0:
												_lowerCAmelCase     :     Union[str, Any]              =   timesteps.astype(dtype=jnp.floataa   )
												_lowerCAmelCase     :     Optional[int]              =   jnp.expand_dims(lowercase__		,						0   )

								_lowerCAmelCase     :     Union[str, Any]              =   self.time_proj(lowercase__   )
								_lowerCAmelCase     :     Any              =   self.time_embedding(lowercase__   )

								# 2. pre-process
								_lowerCAmelCase     :     Tuple              =   jnp.transpose(lowercase__		,						(0, 2, 3, 1)   )
								_lowerCAmelCase     :     int              =   self.conv_in(lowercase__   )

								# 3. down
								_lowerCAmelCase     :     Dict              =   (sample,)
								for down_block in self.down_blocks:
												if isinstance(lowercase__		,						lowercase__   ):
																_lowerCAmelCase     :     Optional[int]              =   down_block(lowercase__		,						lowercase__		,						lowercase__		,						deterministic=not train   )
												else:
																_lowerCAmelCase     :     Union[str, Any]              =   down_block(lowercase__		,						lowercase__		,						deterministic=not train   )
												down_block_res_samples += res_samples

								if down_block_additional_residuals is not None:
												_lowerCAmelCase     :     str              =   ()

												for down_block_res_sample, down_block_additional_residual in zip(
												    lowercase__		,						lowercase__   ):
																down_block_res_sample += down_block_additional_residual
																new_down_block_res_samples += (down_block_res_sample,)

												_lowerCAmelCase     :     str              =   new_down_block_res_samples

								# 4. mid
								_lowerCAmelCase     :     Union[str, Any]              =   self.mid_block(lowercase__		,						lowercase__		,						lowercase__		,						deterministic=not train   )

								if mid_block_additional_residual is not None:
												sample += mid_block_additional_residual

								# 5. up
								for up_block in self.up_blocks:
												_lowerCAmelCase     :     Any              =   down_block_res_samples[-(self.layers_per_block + 1) :]
												_lowerCAmelCase     :     List[Any]              =   down_block_res_samples[: -(self.layers_per_block + 1)]
												if isinstance(lowercase__		,						lowercase__   ):
																_lowerCAmelCase     :     Tuple              =   up_block(
																    lowercase__		,						temb=lowercase__		,						encoder_hidden_states=lowercase__		,						res_hidden_states_tuple=lowercase__		,						deterministic=not train		,						)
												else:
																_lowerCAmelCase     :     Optional[int]              =   up_block(lowercase__		,						temb=lowercase__		,						res_hidden_states_tuple=lowercase__		,						deterministic=not train   )

        # 6. post-process
								_lowerCAmelCase     :     Optional[Any]              =   self.conv_norm_out(lowercase__   )
								_lowerCAmelCase     :     Dict              =   nn.silu(lowercase__   )
								_lowerCAmelCase     :     Tuple              =   self.conv_out(lowercase__   )
								_lowerCAmelCase     :     Optional[Any]              =   jnp.transpose(lowercase__		,						(0, 3, 1, 2)   )

								if not return_dict:
												return (sample,)

								return FlaxUNetaDConditionOutput(sample=lowercase__   )





from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin




@dataclass
class 							lowerCamelCase   (  _UpperCamelCase ):
	_lowerCAmelCase     :		torch.FloatTensor
	_lowerCAmelCase     :		torch.FloatTensor
	_lowerCAmelCase     :		Optional[torch.FloatTensor]										=						None






class 							lowerCamelCase   (  _UpperCamelCase							,						_UpperCamelCase ):
	_lowerCAmelCase     :		Tuple										=						2


	@register_to_config
	def __init__(    self				,     lowercase__ = 0.0_2				,     lowercase__ = 1_0_0				,     lowercase__ = 1.0_0_7				,     lowercase__ = 8_0				,     lowercase__ = 0.0_5				,     lowercase__ = 5_0				,     ):
							# standard deviation of the initial noise distribution
							__UpperCAmelCase			:							Union[str, Any]	    =	sigma_max

							# setable values
							__UpperCAmelCase			:							int	    =	None
							__UpperCAmelCase			:							np.IntTensor	    =	None
							__UpperCAmelCase			:							torch.FloatTensor	    =	None  # sigma(t_i)


	def  A(    self				,     lowercase__				,     lowercase__ = None):
							return sample


	def  A(    self				,     lowercase__				,     lowercase__ = None):
							__UpperCAmelCase			:							int	    =	num_inference_steps
							__UpperCAmelCase			:							List[str]	    =	np.arange(0				,     self.num_inference_steps)[::-1].copy()
							__UpperCAmelCase			:							Any	    =	torch.from_numpy(lowercase__).to(lowercase__)
							__UpperCAmelCase			:							Union[str, Any]	    =	[
							    (
							        self.config.sigma_max**2
							        * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
							    )
							    for i in self.timesteps
							]
							__UpperCAmelCase			:							Tuple	    =	torch.tensor(lowercase__				,     dtype=torch.floataa				,     device=lowercase__)


	def  A(    self				,     lowercase__				,     lowercase__				,     lowercase__ = None):
							if self.config.s_min <= sigma <= self.config.s_max:
													__UpperCAmelCase			:							int	    =	min(self.config.s_churn / self.num_inference_steps				,     2**0.5 - 1)
							else:
													__UpperCAmelCase			:							int	    =	0

							# sample eps ~ N(0, S_noise^2 * I)
							__UpperCAmelCase			:							List[str]	    =	self.config.s_noise * randn_tensor(sample.shape				,     generator=lowercase__).to(sample.device)
							__UpperCAmelCase			:							Optional[int]	    =	sigma + gamma * sigma
							__UpperCAmelCase			:							List[Any]	    =	sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps)

							return sample_hat, sigma_hat


	def  A(    self				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__ = True				,     ):
							__UpperCAmelCase			:							str	    =	sample_hat + sigma_hat * model_output
							__UpperCAmelCase			:							Tuple	    =	(sample_hat - pred_original_sample) / sigma_hat
							__UpperCAmelCase			:							str	    =	sample_hat + (sigma_prev - sigma_hat) * derivative

							if not return_dict:
													return (sample_prev, derivative)

							return KarrasVeOutput(
							    prev_sample=lowercase__				,     derivative=lowercase__				,     pred_original_sample=lowercase__)


	def  A(    self				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__ = True				,     ):
							__UpperCAmelCase			:							Any	    =	sample_prev + sigma_prev * model_output
							__UpperCAmelCase			:							List[str]	    =	(sample_prev - pred_original_sample) / sigma_prev
							__UpperCAmelCase			:							Union[str, Any]	    =	sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr)

							if not return_dict:
													return (sample_prev, derivative)

							return KarrasVeOutput(
							    prev_sample=lowercase__				,     derivative=lowercase__				,     pred_original_sample=lowercase__)


	def  A(    self				,     lowercase__				,     lowercase__				,     lowercase__):
							raise NotImplementedError()

