






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      SCREAMING_SNAKE_CASE  (				__snake_case						):





	"""simple docstring"""




	_A       :				jnp.ndarray
	_A       :				jnp.ndarray






class      SCREAMING_SNAKE_CASE  (				nn.Module						):





	"""simple docstring"""




	_A       :				int
	_A       :				Tuple[int]    =				(16, 32, 96, 256)
	_A       :				jnp.dtype    =				jnp.floataa





	def    lowerCamelCase(self							):
		A_		:						List[str]	    =   nn.Conv(
		    self.block_out_channels[0]     ,     kernel_size=(3, 3)     ,     padding=((1, 1), (1, 1))     ,     dtype=self.dtype     ,     )

		A_		:						Optional[int]	    =   []
		for i in range(len(self.block_out_channels							) - 1							):
			A_		:						List[Any]	    =   self.block_out_channels[i]
			A_		:						Tuple	    =   self.block_out_channels[i + 1]
			A_		:						List[str]	    =   nn.Conv(
			    A_     ,     kernel_size=(3, 3)     ,     padding=((1, 1), (1, 1))     ,     dtype=self.dtype     ,     )
			blocks.append(A_							)
			A_		:						List[str]	    =   nn.Conv(
			    A_     ,     kernel_size=(3, 3)     ,     strides=(2, 2)     ,     padding=((1, 1), (1, 1))     ,     dtype=self.dtype     ,     )
			blocks.append(A_							)
		A_		:						List[Any]	    =   blocks

		A_		:						Dict	    =   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     ,     lowerCAmelCase_							):
		A_		:						Dict	    =   self.conv_in(A_							)
		A_		:						Any	    =   nn.silu(A_							)

		for block in self.blocks:
			A_		:						int	    =   block(A_							)
			A_		:						Any	    =   nn.silu(A_							)

		A_		:						List[Any]	    =   self.conv_out(A_							)

		return embedding






@flax_register_to_config
class      SCREAMING_SNAKE_CASE  (				nn.Module       ,       __snake_case       ,       __snake_case						):





	"""simple docstring"""




	_A       :				int    =				32
	_A       :				int    =				4
	_A       :				Tuple[str]    =				(
	    "CrossAttnDownBlock2D",
	    "CrossAttnDownBlock2D",
	    "CrossAttnDownBlock2D",
	    "DownBlock2D",
	)
	_A       :				Union[bool, Tuple[bool]]    =				False
	_A       :				Tuple[int]    =				(320, 640, 1280, 1280)
	_A       :				int    =				2
	_A       :				Union[int, Tuple[int]]    =				8
	_A       :				Optional[Union[int, Tuple[int]]]    =				None
	_A       :				int    =				1280
	_A       :				float    =				0.0
	_A       :				bool    =				False
	_A       :				jnp.dtype    =				jnp.floataa
	_A       :				bool    =				True
	_A       :				int    =				0
	_A       :				str    =				"rgb"
	_A       :				Tuple[int]    =				(16, 32, 96, 256)





	def    lowerCamelCase(self     ,     lowerCAmelCase_							):
		A_		:						Dict	    =   (1, self.in_channels, self.sample_size, self.sample_size)
		A_		:						Any	    =   jnp.zeros(A_     ,     dtype=jnp.floataa							)
		A_		:						Union[str, Any]	    =   jnp.ones((1,)     ,     dtype=jnp.intaa							)
		A_		:						Dict	    =   jnp.zeros((1, 1, self.cross_attention_dim)     ,     dtype=jnp.floataa							)
		A_		:						List[str]	    =   (1, 3, self.sample_size * 8, self.sample_size * 8)
		A_		:						Optional[Any]	    =   jnp.zeros(A_     ,     dtype=jnp.floataa							)

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

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





	def    lowerCamelCase(self							):
		A_		:						List[str]	    =   self.block_out_channels
		A_		:						Optional[Any]	    =   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_		:						Dict	    =   self.num_attention_heads or self.attention_head_dim

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

		# time
		A_		:						Optional[Any]	    =   FlaxTimesteps(
		    block_out_channels[0]     ,     flip_sin_to_cos=self.flip_sin_to_cos     ,     freq_shift=self.config.freq_shift							)
		A_		:						int	    =   FlaxTimestepEmbedding(A_     ,     dtype=self.dtype							)

		A_		:						int	    =   FlaxControlNetConditioningEmbedding(
		    conditioning_embedding_channels=block_out_channels[0]     ,     block_out_channels=self.conditioning_embedding_out_channels     ,     )

		A_		:						Tuple	    =   self.only_cross_attention
		if isinstance(A_     ,     A_							):
			A_		:						Dict	    =   (only_cross_attention,) * len(self.down_block_types							)

		if isinstance(A_     ,     A_							):
			A_		:						Any	    =   (num_attention_heads,) * len(self.down_block_types							)

		# down
		A_		:						str	    =   []
		A_		:						Optional[Any]	    =   []

		A_		:						int	    =   block_out_channels[0]

		A_		:						str	    =   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							):
			A_		:						Union[str, Any]	    =   output_channel
			A_		:						Optional[Any]	    =   block_out_channels[i]
			A_		:						int	    =   i == len(A_							) - 1

			if down_block_type == "CrossAttnDownBlock2D":
				A_		:						List[Any]	    =   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:
				A_		:						Any	    =   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							):
				A_		:						List[Any]	    =   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:
				A_		:						int	    =   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_							)

		A_		:						List[str]	    =   down_blocks
		A_		:						List[str]	    =   controlnet_down_blocks

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

		A_		:						Dict	    =   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     ,     lowerCAmelCase_     ,     lowerCAmelCase_     ,     lowerCAmelCase_     ,     lowerCAmelCase_     ,     lowerCAmelCase_ = 1.0     ,     lowerCAmelCase_ = True     ,     lowerCAmelCase_ = False     ,     ):
		A_		:						Tuple	    =   self.controlnet_conditioning_channel_order
		if channel_order == "bgr":
			A_		:						str	    =   jnp.flip(A_     ,     axis=1							)

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

		A_		:						int	    =   self.time_proj(A_							)
		A_		:						List[Any]	    =   self.time_embedding(A_							)

		# 2. pre-process
		A_		:						int	    =   jnp.transpose(A_     ,     (0, 2, 3, 1)							)
		A_		:						Dict	    =   self.conv_in(A_							)

		A_		:						List[Any]	    =   jnp.transpose(A_     ,     (0, 2, 3, 1)							)
		A_		:						Optional[int]	    =   self.controlnet_cond_embedding(A_							)
		sample += controlnet_cond

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

		# 4. mid
		A_		:						Any	    =   self.mid_block(A_     ,     A_     ,     A_     ,     deterministic=not train							)

		# 5. contronet blocks
		A_		:						List[Any]	    =   ()
		for down_block_res_sample, controlnet_block in zip(A_     ,     self.controlnet_down_blocks							):
			A_		:						Tuple	    =   controlnet_block(A_							)
			controlnet_down_block_res_samples += (down_block_res_sample,)

		A_		:						int	    =   controlnet_down_block_res_samples

		A_		:						Optional[int]	    =   self.controlnet_mid_block(A_							)

		# 6. scaling
		A_		:						str	    =   [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 gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()



class      SCREAMING_SNAKE_CASE  (				_SCREAMING_SNAKE_CASE       ,       unittest.TestCase						):





	"""simple docstring"""




	_A       :				int    =				PriorTransformer
	_A       :				Dict    =				"""hidden_states"""





	@property
	def    lowerCamelCase(self							):
		A_		:						Tuple	    =   4
		A_		:						Dict	    =   8
		A_		:						str	    =   7

		A_		:						Dict	    =   floats_tensor((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						str	    =   floats_tensor((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						List[str]	    =   floats_tensor((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	def    lowerCamelCase(self     ,     lowerCAmelCase_=0							):
		torch.manual_seed(lowerCAmelCase_							)
		A_		:						Optional[Any]	    =   4
		A_		:						Tuple	    =   8
		A_		:						Tuple	    =   7

		A_		:						int	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						List[str]	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						Tuple	    =   torch.randn((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	@property
	def    lowerCamelCase(self							):
		return (4, 8)





	@property
	def    lowerCamelCase(self							):
		return (4, 8)





	def    lowerCamelCase(self							):
		A_		:						int	    =   {
		    """num_attention_heads""": 2,
		    """attention_head_dim""": 4,
		    """num_layers""": 2,
		    """embedding_dim""": 8,
		    """num_embeddings""": 7,
		    """additional_embeddings""": 4,
		}
		A_		:						Dict	    =   self.dummy_input
		return init_dict, inputs_dict





	def    lowerCamelCase(self							):
		A_    ,     A_		:						Optional[int]	    =   PriorTransformer.from_pretrained(
		    """hf-internal-testing/prior-dummy"""     ,     output_loading_info=lowerCAmelCase_							)
		self.assertIsNotNone(lowerCAmelCase_							)
		self.assertEqual(len(loading_info["""missing_keys"""]							)     ,     0							)

		model.to(lowerCAmelCase_							)
		A_		:						List[Any]	    =   model(**self.dummy_input							)[0]

		assert hidden_states is not None, "Make sure output is not None"





	def    lowerCamelCase(self							):
		A_    ,     A_		:						Union[str, Any]	    =   self.prepare_init_args_and_inputs_for_common()

		A_		:						Union[str, Any]	    =   self.model_class(**lowerCAmelCase_							)
		A_		:						Tuple	    =   inspect.signature(model.forward							)
		# signature.parameters is an OrderedDict => so arg_names order is deterministic
		A_		:						Tuple	    =   [*signature.parameters.keys()]

		A_		:						List[str]	    =   ["""hidden_states""", """timestep"""]
		self.assertListEqual(arg_names[:2]     ,     lowerCAmelCase_							)





	def    lowerCamelCase(self							):
		A_		:						int	    =   PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy"""							)
		A_		:						Optional[Any]	    =   model.to(lowerCAmelCase_							)

		if hasattr(lowerCAmelCase_     ,     """set_default_attn_processor"""							):
			model.set_default_attn_processor()

		A_		:						int	    =   self.get_dummy_seed_input()

		with torch.no_grad():
			A_		:						int	    =   model(**lowerCAmelCase_							)[0]

		A_		:						Any	    =   output[0, :5].flatten().cpu()
		print(lowerCAmelCase_							)

		# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
		# the expected output slices are not the same for CPU and GPU.
		A_		:						Optional[int]	    =   torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239]							)
		self.assertTrue(torch_all_close(lowerCAmelCase_     ,     lowerCAmelCase_     ,     rtol=1e-2							)							)






@slow
class      SCREAMING_SNAKE_CASE  (				unittest.TestCase						):





	"""simple docstring"""




	def    lowerCamelCase(self     ,     lowerCAmelCase_=1     ,     lowerCAmelCase_=768     ,     lowerCAmelCase_=77     ,     lowerCAmelCase_=0							):
		torch.manual_seed(lowerCAmelCase_							)
		A_		:						Tuple	    =   batch_size
		A_		:						Any	    =   embedding_dim
		A_		:						Tuple	    =   num_embeddings

		A_		:						Any	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						Union[str, Any]	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						Tuple	    =   torch.randn((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	def    lowerCamelCase(self							):
		# clean up the VRAM after each test
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	@parameterized.expand(
	    [
	        # fmt: off
	        [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
	        [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
	        # fmt: on
	    ]							)
	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_							):
		A_		:						List[Any]	    =   PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""     ,     subfolder="""prior"""							)
		model.to(lowerCAmelCase_							)
		A_		:						int	    =   self.get_dummy_seed_input(seed=lowerCAmelCase_							)

		with torch.no_grad():
			A_		:						Tuple	    =   model(**lowerCAmelCase_							)[0]

		assert list(sample.shape							) == [1, 768]

		A_		:						Tuple	    =   sample[0, :8].flatten().cpu()
		print(lowerCAmelCase_							)
		A_		:						Optional[int]	    =   torch.tensor(lowerCAmelCase_							)

		assert torch_all_close(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1e-3							)

