




'''simple docstring'''




import functools
import math

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




def 						snake_case__    ( lowerCamelCase__				:   Dict			,		lowerCamelCase__				:   Tuple			,		lowerCamelCase__				:   Any			,		lowerCamelCase__				:   Dict			,		lowerCamelCase__				:   int = 4_0_9_6			)   ->   Dict:
					A_ :     Optional[Any] =							key.shape[-3:]
					A_ :     Optional[int] =							value.shape[-1]
					A_ :     List[str] =							min(lowerCamelCase__			,		lowerCamelCase__			)
					A_ :     int =							query / jnp.sqrt(lowerCamelCase__			)

					@functools.partial(jax.checkpoint			,		prevent_cse=lowerCamelCase__			)
					def summarize_chunk(lowerCamelCase__				:   int			,		lowerCamelCase__				:   int			,		lowerCamelCase__				:   Any			):
										A_ :     Dict =							jnp.einsum('''...qhd,...khd->...qhk'''			,		lowerCamelCase__			,		lowerCamelCase__			,		precision=lowerCamelCase__			)

										A_ :     str =							jnp.max(lowerCamelCase__			,		axis=-1			,		keepdims=lowerCamelCase__			)
										A_ :     Dict =							jax.lax.stop_gradient(lowerCamelCase__			)
										A_ :     Optional[Any] =							jnp.exp(attn_weights - max_score			)

										A_ :     Optional[Any] =							jnp.einsum('''...vhf,...qhv->...qhf'''			,		lowerCamelCase__			,		lowerCamelCase__			,		precision=lowerCamelCase__			)
										A_ :     Optional[Any] =							jnp.einsum('''...qhk->...qh'''			,		lowerCamelCase__			)

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

					def chunk_scanner(lowerCamelCase__				:   Optional[Any]			):
										# julienne key array
										A_ :     List[str] =							jax.lax.dynamic_slice(
										    operand=lowerCamelCase__			,		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
										A_ :     int =							jax.lax.dynamic_slice(
										    operand=lowerCamelCase__			,		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(lowerCamelCase__			,		lowerCamelCase__			,		lowerCamelCase__			)

					A_ :     List[Any] =							jax.lax.map(f=lowerCamelCase__			,		xs=jnp.arange(0			,		lowerCamelCase__			,		lowerCamelCase__			)			)

					A_ :     Union[str, Any] =							jnp.max(lowerCamelCase__			,		axis=0			,		keepdims=lowerCamelCase__			)
					A_ :     str =							jnp.exp(chunk_max - global_max			)

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

					A_ :     Union[str, Any] =							chunk_values.sum(axis=0			)
					A_ :     Optional[Any] =							jnp.expand_dims(lowerCamelCase__			,		-1			).sum(axis=0			)

					return all_values / all_weights




def 						snake_case__    ( lowerCamelCase__				:   Any			,		lowerCamelCase__				:   List[str]			,		lowerCamelCase__				:   List[str]			,		lowerCamelCase__				:   Optional[int]=jax.lax.Precision.HIGHEST			,		lowerCamelCase__				:   int = 1_0_2_4			,		lowerCamelCase__				:   int = 4_0_9_6			)   ->   str:
					A_ :     Any =							query.shape[-3:]

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

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

					A_ :     List[str] =							jax.lax.scan(
					    f=lowerCamelCase__			,		init=0			,		xs=lowerCamelCase__			,		length=math.ceil(num_q / query_chunk_size			)			)

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






class     UpperCamelCase_     (nn.Module    ):







						"""simple docstring"""


						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					8
						_lowerCAmelCase    						=					6_4
						_lowerCAmelCase    						=					0.0
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					jnp.floataa







						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							self.dim_head * self.heads
											A_ :     Tuple =							self.dim_head**-0.5

											# Weights were exported with old names {to_q, to_k, to_v, to_out}
											A_ :     List[Any] =							nn.Dense(_lowerCamelCase							, use_bias=_lowerCamelCase							, dtype=self.dtype							, name='''to_q'''    )
											A_ :     List[Any] =							nn.Dense(_lowerCamelCase							, use_bias=_lowerCamelCase							, dtype=self.dtype							, name='''to_k'''    )
											A_ :     str =							nn.Dense(_lowerCamelCase							, use_bias=_lowerCamelCase							, dtype=self.dtype							, name='''to_v'''    )

											A_ :     List[Any] =							nn.Dense(self.query_dim							, dtype=self.dtype							, name='''to_out_0'''    )
											A_ :     Optional[int] =							nn.Dropout(rate=self.dropout    )







						def       _a    (     self      :       Optional[int]							, _lowerCamelCase      :       List[Any]    ):





											"""simple docstring"""







											A_ :     Optional[Any] =							tensor.shape
											A_ :     Dict =							self.heads
											A_ :     Tuple =							tensor.reshape(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, dim // head_size    )
											A_ :     Optional[Any] =							jnp.transpose(_lowerCamelCase							, (0, 2, 1, 3)    )
											A_ :     List[Any] =							tensor.reshape(batch_size * head_size							, _lowerCamelCase							, dim // head_size    )
											return tensor







						def       _a    (     self      :       Any							, _lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ :     Any =							tensor.shape
											A_ :     Any =							self.heads
											A_ :     Optional[int] =							tensor.reshape(batch_size // head_size							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )
											A_ :     List[str] =							jnp.transpose(_lowerCamelCase							, (0, 2, 1, 3)    )
											A_ :     Optional[int] =							tensor.reshape(batch_size // head_size							, _lowerCamelCase							, dim * head_size    )
											return tensor







						def __call__(     self      :       Optional[Any]							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       Tuple=None							, _lowerCamelCase      :       int=True    ):





											"""simple docstring"""







											A_ :     int =							hidden_states if context is None else context

											A_ :     List[str] =							self.query(_lowerCamelCase    )
											A_ :     Any =							self.key(_lowerCamelCase    )
											A_ :     Optional[Any] =							self.value(_lowerCamelCase    )

											A_ :     Union[str, Any] =							self.reshape_heads_to_batch_dim(_lowerCamelCase    )
											A_ :     List[str] =							self.reshape_heads_to_batch_dim(_lowerCamelCase    )
											A_ :     Union[str, Any] =							self.reshape_heads_to_batch_dim(_lowerCamelCase    )

											if self.use_memory_efficient_attention:
																A_ :     str =							query_states.transpose(1							, 0							, 2    )
																A_ :     Optional[int] =							key_states.transpose(1							, 0							, 2    )
																A_ :     int =							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

																A_ :     Any =							query_states.shape[-3]
																if flatten_latent_dim % 64 == 0:
																					A_ :     Optional[Any] =							int(flatten_latent_dim / 64    )
																elif flatten_latent_dim % 16 == 0:
																					A_ :     str =							int(flatten_latent_dim / 16    )
																elif flatten_latent_dim % 4 == 0:
																					A_ :     Any =							int(flatten_latent_dim / 4    )
																else:
																					A_ :     Any =							int(_lowerCamelCase    )

																A_ :     Any =							jax_memory_efficient_attention(
																    _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, query_chunk_size=_lowerCamelCase							, key_chunk_size=4096 * 4    )

																A_ :     Optional[Any] =							hidden_states.transpose(1							, 0							, 2    )
											else:
																# compute attentions
																A_ :     Tuple =							jnp.einsum('''b i d, b j d->b i j'''							, _lowerCamelCase							, _lowerCamelCase    )
																A_ :     Optional[Any] =							attention_scores * self.scale
																A_ :     Union[str, Any] =							nn.softmax(_lowerCamelCase							, axis=2    )

																# attend to values
																A_ :     Dict =							jnp.einsum('''b i j, b j d -> b i d'''							, _lowerCamelCase							, _lowerCamelCase    )

											A_ :     Tuple =							self.reshape_batch_dim_to_heads(_lowerCamelCase    )
											A_ :     Union[str, Any] =							self.proj_attn(_lowerCamelCase    )
											return self.dropout_layer(_lowerCamelCase							, deterministic=_lowerCamelCase    )






class     UpperCamelCase_     (nn.Module    ):







						"""simple docstring"""


						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					0.0
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					jnp.floataa
						_lowerCAmelCase    						=					False







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ :     int =							FlaxAttention(
											    self.dim							, self.n_heads							, self.d_head							, self.dropout							, self.use_memory_efficient_attention							, dtype=self.dtype    )
											# cross attention
											A_ :     List[str] =							FlaxAttention(
											    self.dim							, self.n_heads							, self.d_head							, self.dropout							, self.use_memory_efficient_attention							, dtype=self.dtype    )
											A_ :     Union[str, Any] =							FlaxFeedForward(dim=self.dim							, dropout=self.dropout							, dtype=self.dtype    )
											A_ :     Tuple =							nn.LayerNorm(epsilon=1E-5							, dtype=self.dtype    )
											A_ :     Dict =							nn.LayerNorm(epsilon=1E-5							, dtype=self.dtype    )
											A_ :     Dict =							nn.LayerNorm(epsilon=1E-5							, dtype=self.dtype    )
											A_ :     List[Any] =							nn.Dropout(rate=self.dropout    )







						def __call__(     self      :       Optional[Any]							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Optional[int]=True    ):





											"""simple docstring"""







											A_ :     Optional[int] =							hidden_states
											if self.only_cross_attention:
																A_ :     Optional[Any] =							self.attna(self.norma(_lowerCamelCase    )							, _lowerCamelCase							, deterministic=_lowerCamelCase    )
											else:
																A_ :     Optional[int] =							self.attna(self.norma(_lowerCamelCase    )							, deterministic=_lowerCamelCase    )
											A_ :     List[str] =							hidden_states + residual

											# cross attention
											A_ :     Optional[int] =							hidden_states
											A_ :     int =							self.attna(self.norma(_lowerCamelCase    )							, _lowerCamelCase							, deterministic=_lowerCamelCase    )
											A_ :     List[str] =							hidden_states + residual

											# feed forward
											A_ :     str =							hidden_states
											A_ :     Dict =							self.ff(self.norma(_lowerCamelCase    )							, deterministic=_lowerCamelCase    )
											A_ :     List[str] =							hidden_states + residual

											return self.dropout_layer(_lowerCamelCase							, deterministic=_lowerCamelCase    )






class     UpperCamelCase_     (nn.Module    ):







						"""simple docstring"""


						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					1
						_lowerCAmelCase    						=					0.0
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					jnp.floataa
						_lowerCAmelCase    						=					False







						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											A_ :     int =							nn.GroupNorm(num_groups=32							, epsilon=1E-5    )

											A_ :     List[str] =							self.n_heads * self.d_head
											if self.use_linear_projection:
																A_ :     Optional[Any] =							nn.Dense(_lowerCamelCase							, dtype=self.dtype    )
											else:
																A_ :     Dict =							nn.Conv(
																    _lowerCamelCase							, kernel_size=(1, 1)							, strides=(1, 1)							, padding='''VALID'''							, dtype=self.dtype							, )

											A_ :     Union[str, Any] =							[
											    FlaxBasicTransformerBlock(
											        _lowerCamelCase							, 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:
																A_ :     Optional[int] =							nn.Dense(_lowerCamelCase							, dtype=self.dtype    )
											else:
																A_ :     str =							nn.Conv(
																    _lowerCamelCase							, kernel_size=(1, 1)							, strides=(1, 1)							, padding='''VALID'''							, dtype=self.dtype							, )

											A_ :     int =							nn.Dropout(rate=self.dropout    )







						def __call__(     self      :       Tuple							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       List[str]=True    ):





											"""simple docstring"""







											A_ :     Tuple =							hidden_states.shape
											A_ :     List[str] =							hidden_states
											A_ :     List[str] =							self.norm(_lowerCamelCase    )
											if self.use_linear_projection:
																A_ :     Optional[Any] =							hidden_states.reshape(_lowerCamelCase							, height * width							, _lowerCamelCase    )
																A_ :     Tuple =							self.proj_in(_lowerCamelCase    )
											else:
																A_ :     int =							self.proj_in(_lowerCamelCase    )
																A_ :     Optional[int] =							hidden_states.reshape(_lowerCamelCase							, height * width							, _lowerCamelCase    )

											for transformer_block in self.transformer_blocks:
																A_ :     Dict =							transformer_block(_lowerCamelCase							, _lowerCamelCase							, deterministic=_lowerCamelCase    )

											if self.use_linear_projection:
																A_ :     Any =							self.proj_out(_lowerCamelCase    )
																A_ :     Union[str, Any] =							hidden_states.reshape(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )
											else:
																A_ :     Optional[Any] =							hidden_states.reshape(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )
																A_ :     List[Any] =							self.proj_out(_lowerCamelCase    )

											A_ :     Any =							hidden_states + residual
											return self.dropout_layer(_lowerCamelCase							, deterministic=_lowerCamelCase    )






class     UpperCamelCase_     (nn.Module    ):







						"""simple docstring"""


						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					0.0
						_lowerCAmelCase    						=					jnp.floataa







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											A_ :     List[str] =							FlaxGEGLU(self.dim							, self.dropout							, self.dtype    )
											A_ :     Optional[int] =							nn.Dense(self.dim							, dtype=self.dtype    )







						def __call__(     self      :       Dict							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Union[str, Any]=True    ):





											"""simple docstring"""







											A_ :     Optional[Any] =							self.net_a(_lowerCamelCase							, deterministic=_lowerCamelCase    )
											A_ :     Dict =							self.net_a(_lowerCamelCase    )
											return hidden_states






class     UpperCamelCase_     (nn.Module    ):







						"""simple docstring"""


						_lowerCAmelCase    						=					4_2
						_lowerCAmelCase    						=					0.0
						_lowerCAmelCase    						=					jnp.floataa







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ :     int =							self.dim * 4
											A_ :     List[Any] =							nn.Dense(inner_dim * 2							, dtype=self.dtype    )
											A_ :     List[Any] =							nn.Dropout(rate=self.dropout    )







						def __call__(     self      :       int							, _lowerCamelCase      :       Optional[Any]							, _lowerCamelCase      :       Any=True    ):





											"""simple docstring"""







											A_ :     str =							self.proj(_lowerCamelCase    )
											A_ :     str =							jnp.split(_lowerCamelCase							, 2							, axis=2    )
											return self.dropout_layer(hidden_linear * nn.gelu(_lowerCamelCase    )							, deterministic=_lowerCamelCase    )









'''simple docstring'''




import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()






class     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						@property
						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ :     Dict =							1
											A_ :     Any =							3
											A_ :     Any =							(32, 32)

											A_ :     Dict =							floats_tensor((batch_size, num_channels) + sizes							, rng=random.Random(0    )    ).to(_lowerCamelCase    )
											return image







						@property
						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											torch.manual_seed(0    )
											A_ :     Optional[int] =							UNetaDModel(
											    block_out_channels=(32, 64)							, layers_per_block=2							, sample_size=32							, in_channels=6							, out_channels=3							, down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')							, up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')							, )
											return model







						@property
						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											torch.manual_seed(0    )
											A_ :     Dict =							VQModel(
											    block_out_channels=[32, 64]							, in_channels=3							, out_channels=3							, down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']							, up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']							, latent_channels=3							, )
											return model







						def       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ :     Any =							'''cpu'''
											A_ :     Optional[Any] =							self.dummy_uncond_unet
											A_ :     str =							DDIMScheduler()
											A_ :     Optional[int] =							self.dummy_vq_model

											A_ :     Optional[Any] =							LDMSuperResolutionPipeline(unet=_lowerCamelCase							, vqvae=_lowerCamelCase							, scheduler=_lowerCamelCase    )
											ldm.to(_lowerCamelCase    )
											ldm.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     Any =							self.dummy_image.to(_lowerCamelCase    )

											A_ :     str =							torch.Generator(device=_lowerCamelCase    ).manual_seed(0    )
											A_ :     List[Any] =							ldm(image=_lowerCamelCase							, generator=_lowerCamelCase							, num_inference_steps=2							, output_type='''numpy'''    ).images

											A_ :     str =							image[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											A_ :     List[Any] =							np.array([0.86_78, 0.82_45, 0.63_81, 0.68_30, 0.43_85, 0.55_99, 0.46_41, 0.62_01, 0.51_50]    )

											assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2







						@unittest.skipIf(torch_device != '''cuda'''							, '''This test requires a GPU'''    )
						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Tuple =							self.dummy_uncond_unet
											A_ :     List[str] =							DDIMScheduler()
											A_ :     str =							self.dummy_vq_model

											# put models in fp16
											A_ :     List[Any] =							unet.half()
											A_ :     Optional[int] =							vqvae.half()

											A_ :     Optional[Any] =							LDMSuperResolutionPipeline(unet=_lowerCamelCase							, vqvae=_lowerCamelCase							, scheduler=_lowerCamelCase    )
											ldm.to(_lowerCamelCase    )
											ldm.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     Optional[Any] =							self.dummy_image.to(_lowerCamelCase    )

											A_ :     Optional[Any] =							ldm(_lowerCamelCase							, num_inference_steps=2							, output_type='''numpy'''    ).images

											assert image.shape == (1, 64, 64, 3)






@slow
@require_torch
class     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ :     Optional[int] =							load_image(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
											    '''/vq_diffusion/teddy_bear_pool.png'''    )
											A_ :     List[str] =							init_image.resize((64, 64)							, resample=PIL_INTERPOLATION['''lanczos''']    )

											A_ :     List[str] =							LDMSuperResolutionPipeline.from_pretrained('''duongna/ldm-super-resolution'''							, device_map='''auto'''    )
											ldm.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     List[str] =							torch.manual_seed(0    )
											A_ :     str =							ldm(image=_lowerCamelCase							, generator=_lowerCamelCase							, num_inference_steps=20							, output_type='''numpy'''    ).images

											A_ :     str =							image[0, -3:, -3:, -1]

											assert image.shape == (1, 256, 256, 3)
											A_ :     Optional[Any] =							np.array([0.76_44, 0.76_79, 0.76_42, 0.76_33, 0.76_66, 0.75_60, 0.74_25, 0.72_57, 0.69_07]    )

											assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2



