



import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp


def 		_UpperCAmelCase    (				snake_case     ,						snake_case     ,						snake_case     ,						snake_case     ,						snake_case = 40_96      ):





					"""simple docstring"""

					_lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase      =   key.shape[-3:]
					_lowerCAmelCase      =   value.shape[-1]
					_lowerCAmelCase      =   min(snake_case     ,						snake_case      )
					_lowerCAmelCase      =   query / jnp.sqrt(snake_case      )

					@functools.partial(jax.checkpoint     ,						prevent_cse=snake_case      )
					def summarize_chunk(snake_case     ,						snake_case     ,						snake_case      ):
										_lowerCAmelCase      =   jnp.einsum("""...qhd,...khd->...qhk"""     ,						snake_case     ,						snake_case     ,						precision=snake_case      )

										_lowerCAmelCase      =   jnp.max(snake_case     ,						axis=-1     ,						keepdims=snake_case      )
										_lowerCAmelCase      =   jax.lax.stop_gradient(snake_case      )
										_lowerCAmelCase      =   jnp.exp(attn_weights - max_score      )

										_lowerCAmelCase      =   jnp.einsum("""...vhf,...qhv->...qhf"""     ,						snake_case     ,						snake_case     ,						precision=snake_case      )
										_lowerCAmelCase      =   jnp.einsum("""...qhk->...qh"""     ,						snake_case      )

										return (exp_values, exp_weights.sum(axis=-1      ), max_score)

					def chunk_scanner(snake_case      ):
										# julienne key array
										_lowerCAmelCase      =   jax.lax.dynamic_slice(
										    operand=snake_case     ,						start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]     ,						slice_sizes=list(key.shape[:-3]      ) + [key_chunk_size, num_heads, k_features]     ,						)

										# julienne value array
										_lowerCAmelCase      =   jax.lax.dynamic_slice(
										    operand=snake_case     ,						start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]     ,						slice_sizes=list(value.shape[:-3]      ) + [key_chunk_size, num_heads, v_features]     ,						)

										return summarize_chunk(snake_case     ,						snake_case     ,						snake_case      )

					_lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase      =   jax.lax.map(f=snake_case     ,						xs=jnp.arange(0     ,						snake_case     ,						snake_case      )      )

					_lowerCAmelCase      =   jnp.max(snake_case     ,						axis=0     ,						keepdims=snake_case      )
					_lowerCAmelCase      =   jnp.exp(chunk_max - global_max      )

					chunk_values *= jnp.expand_dims(snake_case     ,						axis=-1      )
					chunk_weights *= max_diffs

					_lowerCAmelCase      =   chunk_values.sum(axis=0      )
					_lowerCAmelCase      =   jnp.expand_dims(snake_case     ,						-1      ).sum(axis=0      )

					return all_values / all_weights


def 		_UpperCAmelCase    (				snake_case     ,						snake_case     ,						snake_case     ,						snake_case=jax.lax.Precision.HIGHEST     ,						snake_case = 10_24     ,						snake_case = 40_96      ):





					"""simple docstring"""

					_lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase      =   query.shape[-3:]

					def chunk_scanner(snake_case     ,						snake_case      ):
										# julienne query array
										_lowerCAmelCase      =   jax.lax.dynamic_slice(
										    operand=snake_case     ,						start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]     ,						slice_sizes=list(query.shape[:-3]      ) + [min(snake_case     ,						snake_case      ), num_heads, q_features]     ,						)

										return (
										    chunk_idx + query_chunk_size,  # unused ignore it
										    _query_chunk_attention(
										        query=snake_case     ,						key=snake_case     ,						value=snake_case     ,						precision=snake_case     ,						key_chunk_size=snake_case      ),
										)

					_lowerCAmelCase     ,     _lowerCAmelCase      =   jax.lax.scan(
					    f=snake_case     ,						init=0     ,						xs=snake_case     ,						length=math.ceil(num_q / query_chunk_size      )      )

					return jnp.concatenate(snake_case     ,						axis=-3      )  # fuse the chunked result back





class  __lowerCAmelCase			(					nn.Module	):
				__lowerCamelCase			=    42
				__lowerCamelCase			=    8
				__lowerCamelCase			=    64
				__lowerCamelCase			=    0.0
				__lowerCamelCase			=    False
				__lowerCamelCase			=    jnp.floataa




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.dim_head * self.heads
									_lowerCAmelCase      =   self.dim_head**-0.5

									# Weights were exported with old names {to_q, to_k, to_v, to_out}
									_lowerCAmelCase      =   nn.Dense(_snake_case						,    use_bias=_snake_case						,    dtype=self.dtype						,    name="""to_q"""		)
									_lowerCAmelCase      =   nn.Dense(_snake_case						,    use_bias=_snake_case						,    dtype=self.dtype						,    name="""to_k"""		)
									_lowerCAmelCase      =   nn.Dense(_snake_case						,    use_bias=_snake_case						,    dtype=self.dtype						,    name="""to_v"""		)

									_lowerCAmelCase      =   nn.Dense(self.query_dim						,    dtype=self.dtype						,    name="""to_out_0"""		)
									_lowerCAmelCase      =   nn.Dropout(rate=self.dropout		)




				def 							snake_case       (							self						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase      =   tensor.shape
									_lowerCAmelCase      =   self.heads
									_lowerCAmelCase      =   tensor.reshape(_snake_case						,    _snake_case						,    _snake_case						,    dim // head_size		)
									_lowerCAmelCase      =   jnp.transpose(_snake_case						,    (0, 2, 1, 3)		)
									_lowerCAmelCase      =   tensor.reshape(batch_size * head_size						,    _snake_case						,    dim // head_size		)
									return tensor




				def 							snake_case       (							self						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase      =   tensor.shape
									_lowerCAmelCase      =   self.heads
									_lowerCAmelCase      =   tensor.reshape(batch_size // head_size						,    _snake_case						,    _snake_case						,    _snake_case		)
									_lowerCAmelCase      =   jnp.transpose(_snake_case						,    (0, 2, 1, 3)		)
									_lowerCAmelCase      =   tensor.reshape(batch_size // head_size						,    _snake_case						,    dim * head_size		)
									return tensor




				def __call__(							self						,    _snake_case						,    _snake_case=None						,    _snake_case=True		):



									"""simple docstring"""





									_lowerCAmelCase      =   hidden_states if context is None else context

									_lowerCAmelCase      =   self.query(_snake_case		)
									_lowerCAmelCase      =   self.key(_snake_case		)
									_lowerCAmelCase      =   self.value(_snake_case		)

									_lowerCAmelCase      =   self.reshape_heads_to_batch_dim(_snake_case		)
									_lowerCAmelCase      =   self.reshape_heads_to_batch_dim(_snake_case		)
									_lowerCAmelCase      =   self.reshape_heads_to_batch_dim(_snake_case		)

									if self.use_memory_efficient_attention:
														_lowerCAmelCase      =   query_states.transpose(1						,    0						,    2		)
														_lowerCAmelCase      =   key_states.transpose(1						,    0						,    2		)
														_lowerCAmelCase      =   value_states.transpose(1						,    0						,    2		)

														# this if statement create a chunk size for each layer of the unet
														# the chunk size is equal to the query_length dimension of the deepest layer of the unet

														_lowerCAmelCase      =   query_states.shape[-3]
														if flatten_latent_dim % 64 == 0:
																			_lowerCAmelCase      =   int(flatten_latent_dim / 64		)
														elif flatten_latent_dim % 16 == 0:
																			_lowerCAmelCase      =   int(flatten_latent_dim / 16		)
														elif flatten_latent_dim % 4 == 0:
																			_lowerCAmelCase      =   int(flatten_latent_dim / 4		)
														else:
																			_lowerCAmelCase      =   int(_snake_case		)

														_lowerCAmelCase      =   jax_memory_efficient_attention(
														    _snake_case						,    _snake_case						,    _snake_case						,    query_chunk_size=_snake_case						,    key_chunk_size=4096 * 4		)

														_lowerCAmelCase      =   hidden_states.transpose(1						,    0						,    2		)
									else:
														# compute attentions
														_lowerCAmelCase      =   jnp.einsum("""b i d, b j d->b i j"""						,    _snake_case						,    _snake_case		)
														_lowerCAmelCase      =   attention_scores * self.scale
														_lowerCAmelCase      =   nn.softmax(_snake_case						,    axis=2		)

														# attend to values
														_lowerCAmelCase      =   jnp.einsum("""b i j, b j d -> b i d"""						,    _snake_case						,    _snake_case		)

									_lowerCAmelCase      =   self.reshape_batch_dim_to_heads(_snake_case		)
									_lowerCAmelCase      =   self.proj_attn(_snake_case		)
									return self.dropout_layer(_snake_case						,    deterministic=_snake_case		)





class  __lowerCAmelCase			(					nn.Module	):
				__lowerCamelCase			=    42
				__lowerCamelCase			=    42
				__lowerCamelCase			=    42
				__lowerCamelCase			=    0.0
				__lowerCamelCase			=    False
				__lowerCamelCase			=    jnp.floataa
				__lowerCamelCase			=    False




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   FlaxAttention(
									    self.dim						,    self.n_heads						,    self.d_head						,    self.dropout						,    self.use_memory_efficient_attention						,    dtype=self.dtype		)
									# cross attention
									_lowerCAmelCase      =   FlaxAttention(
									    self.dim						,    self.n_heads						,    self.d_head						,    self.dropout						,    self.use_memory_efficient_attention						,    dtype=self.dtype		)
									_lowerCAmelCase      =   FlaxFeedForward(dim=self.dim						,    dropout=self.dropout						,    dtype=self.dtype		)
									_lowerCAmelCase      =   nn.LayerNorm(epsilon=1e-5						,    dtype=self.dtype		)
									_lowerCAmelCase      =   nn.LayerNorm(epsilon=1e-5						,    dtype=self.dtype		)
									_lowerCAmelCase      =   nn.LayerNorm(epsilon=1e-5						,    dtype=self.dtype		)
									_lowerCAmelCase      =   nn.Dropout(rate=self.dropout		)




				def __call__(							self						,    _snake_case						,    _snake_case						,    _snake_case=True		):



									"""simple docstring"""





									_lowerCAmelCase      =   hidden_states
									if self.only_cross_attention:
														_lowerCAmelCase      =   self.attna(self.norma(_snake_case		)						,    _snake_case						,    deterministic=_snake_case		)
									else:
														_lowerCAmelCase      =   self.attna(self.norma(_snake_case		)						,    deterministic=_snake_case		)
									_lowerCAmelCase      =   hidden_states + residual

									# cross attention
									_lowerCAmelCase      =   hidden_states
									_lowerCAmelCase      =   self.attna(self.norma(_snake_case		)						,    _snake_case						,    deterministic=_snake_case		)
									_lowerCAmelCase      =   hidden_states + residual

									# feed forward
									_lowerCAmelCase      =   hidden_states
									_lowerCAmelCase      =   self.ff(self.norma(_snake_case		)						,    deterministic=_snake_case		)
									_lowerCAmelCase      =   hidden_states + residual

									return self.dropout_layer(_snake_case						,    deterministic=_snake_case		)





class  __lowerCAmelCase			(					nn.Module	):
				__lowerCamelCase			=    42
				__lowerCamelCase			=    42
				__lowerCamelCase			=    42
				__lowerCamelCase			=    1
				__lowerCamelCase			=    0.0
				__lowerCamelCase			=    False
				__lowerCamelCase			=    False
				__lowerCamelCase			=    jnp.floataa
				__lowerCamelCase			=    False




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   nn.GroupNorm(num_groups=32						,    epsilon=1e-5		)

									_lowerCAmelCase      =   self.n_heads * self.d_head
									if self.use_linear_projection:
														_lowerCAmelCase      =   nn.Dense(_snake_case						,    dtype=self.dtype		)
									else:
														_lowerCAmelCase      =   nn.Conv(
														    _snake_case						,    kernel_size=(1, 1)						,    strides=(1, 1)						,    padding="""VALID"""						,    dtype=self.dtype						,    )

									_lowerCAmelCase      =   [
									    FlaxBasicTransformerBlock(
									        _snake_case						,    self.n_heads						,    self.d_head						,    dropout=self.dropout						,    only_cross_attention=self.only_cross_attention						,    dtype=self.dtype						,    use_memory_efficient_attention=self.use_memory_efficient_attention						,    )
									    for _ in range(self.depth		)
									]

									if self.use_linear_projection:
														_lowerCAmelCase      =   nn.Dense(_snake_case						,    dtype=self.dtype		)
									else:
														_lowerCAmelCase      =   nn.Conv(
														    _snake_case						,    kernel_size=(1, 1)						,    strides=(1, 1)						,    padding="""VALID"""						,    dtype=self.dtype						,    )

									_lowerCAmelCase      =   nn.Dropout(rate=self.dropout		)




				def __call__(							self						,    _snake_case						,    _snake_case						,    _snake_case=True		):



									"""simple docstring"""





									_lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase     ,     _lowerCAmelCase      =   hidden_states.shape
									_lowerCAmelCase      =   hidden_states
									_lowerCAmelCase      =   self.norm(_snake_case		)
									if self.use_linear_projection:
														_lowerCAmelCase      =   hidden_states.reshape(_snake_case						,    height * width						,    _snake_case		)
														_lowerCAmelCase      =   self.proj_in(_snake_case		)
									else:
														_lowerCAmelCase      =   self.proj_in(_snake_case		)
														_lowerCAmelCase      =   hidden_states.reshape(_snake_case						,    height * width						,    _snake_case		)

									for transformer_block in self.transformer_blocks:
														_lowerCAmelCase      =   transformer_block(_snake_case						,    _snake_case						,    deterministic=_snake_case		)

									if self.use_linear_projection:
														_lowerCAmelCase      =   self.proj_out(_snake_case		)
														_lowerCAmelCase      =   hidden_states.reshape(_snake_case						,    _snake_case						,    _snake_case						,    _snake_case		)
									else:
														_lowerCAmelCase      =   hidden_states.reshape(_snake_case						,    _snake_case						,    _snake_case						,    _snake_case		)
														_lowerCAmelCase      =   self.proj_out(_snake_case		)

									_lowerCAmelCase      =   hidden_states + residual
									return self.dropout_layer(_snake_case						,    deterministic=_snake_case		)





class  __lowerCAmelCase			(					nn.Module	):
				__lowerCamelCase			=    42
				__lowerCamelCase			=    0.0
				__lowerCamelCase			=    jnp.floataa




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   FlaxGEGLU(self.dim						,    self.dropout						,    self.dtype		)
									_lowerCAmelCase      =   nn.Dense(self.dim						,    dtype=self.dtype		)




				def __call__(							self						,    _snake_case						,    _snake_case=True		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.net_a(_snake_case						,    deterministic=_snake_case		)
									_lowerCAmelCase      =   self.net_a(_snake_case		)
									return hidden_states





class  __lowerCAmelCase			(					nn.Module	):
				__lowerCamelCase			=    42
				__lowerCamelCase			=    0.0
				__lowerCamelCase			=    jnp.floataa




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.dim * 4
									_lowerCAmelCase      =   nn.Dense(inner_dim * 2						,    dtype=self.dtype		)
									_lowerCAmelCase      =   nn.Dropout(rate=self.dropout		)




				def __call__(							self						,    _snake_case						,    _snake_case=True		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.proj(_snake_case		)
									_lowerCAmelCase     ,     _lowerCAmelCase      =   jnp.split(_snake_case						,    2						,    axis=2		)
									return self.dropout_layer(hidden_linear * nn.gelu(_snake_case		)						,    deterministic=_snake_case		)








import os

from packaging import version

from .. import __version__
from .accelerate_utils import apply_forward_hook
from .constants import (
    CONFIG_NAME,
    DEPRECATED_REVISION_ARGS,
    DIFFUSERS_CACHE,
    DIFFUSERS_DYNAMIC_MODULE_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    ONNX_EXTERNAL_WEIGHTS_NAME,
    ONNX_WEIGHTS_NAME,
    SAFETENSORS_WEIGHTS_NAME,
    TEXT_ENCODER_ATTN_MODULE,
    WEIGHTS_NAME,
)
from .deprecation_utils import deprecate
from .doc_utils import replace_example_docstring
from .dynamic_modules_utils import get_class_from_dynamic_module
from .hub_utils import (
    HF_HUB_OFFLINE,
    _add_variant,
    _get_model_file,
    extract_commit_hash,
    http_user_agent,
)
from .import_utils import (
    BACKENDS_MAPPING,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    is_accelerate_available,
    is_accelerate_version,
    is_bsa_available,
    is_flax_available,
    is_ftfy_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_omegaconf_available,
    is_onnx_available,
    is_safetensors_available,
    is_scipy_available,
    is_tensorboard_available,
    is_tf_available,
    is_torch_available,
    is_torch_version,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    is_wandb_available,
    is_xformers_available,
    requires_backends,
)
from .logging import get_logger
from .outputs import BaseOutput
from .pil_utils import PIL_INTERPOLATION, numpy_to_pil, pt_to_pil
from .torch_utils import is_compiled_module, randn_tensor


if is_torch_available():
	from .testing_utils import (
	    floats_tensor,
	    load_hf_numpy,
	    load_image,
	    load_numpy,
	    load_pt,
	    nightly,
	    parse_flag_from_env,
	    print_tensor_test,
	    require_torch_a,
	    require_torch_gpu,
	    skip_mps,
	    slow,
	    torch_all_close,
	    torch_device,
	)
	from .torch_utils import maybe_allow_in_graph

from .testing_utils import export_to_gif, export_to_video


A__     				=	get_logger(__name__)


def 		_UpperCAmelCase    (				snake_case      ):





					"""simple docstring"""

					if version.parse(snake_case      ) < version.parse(snake_case      ):
										if "dev" in min_version:
															_lowerCAmelCase      =   (
															    """This example requires a source install from HuggingFace diffusers (see """
															    """`https://huggingface.co/docs/diffusers/installation#install-from-source`),"""
															)
										else:
															_lowerCAmelCase      =   F'This example requires a minimum version of {min_version},'
										error_message += F' but the version found is {__version__}.\n'
										raise ImportError(snake_case      )



