




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        UpperCamelCase_					(				UpperCAmelCase__  ):

				'''simple docstring'''






				UpperCAmelCase__													=   42


@flax_register_to_config
class        UpperCamelCase_					(				nn.Module	,      UpperCAmelCase__	,      UpperCAmelCase__  ):

				'''simple docstring'''






				UpperCAmelCase__													=   32
				UpperCAmelCase__													=   4
				UpperCAmelCase__													=   4
				UpperCAmelCase__													=   (
				    '''CrossAttnDownBlock2D''',
				    '''CrossAttnDownBlock2D''',
				    '''CrossAttnDownBlock2D''',
				    '''DownBlock2D''',
				)
				UpperCAmelCase__													=   ('''UpBlock2D''', '''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''')
				UpperCAmelCase__													=   False
				UpperCAmelCase__													=   (320, 640, 1280, 1280)
				UpperCAmelCase__													=   2
				UpperCAmelCase__													=   8
				UpperCAmelCase__													=   None
				UpperCAmelCase__													=   1280
				UpperCAmelCase__													=   0.0
				UpperCAmelCase__													=   False
				UpperCAmelCase__													=   jnp.floataa
				UpperCAmelCase__													=   True
				UpperCAmelCase__													=   0
				UpperCAmelCase__													=   False


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





											'''simple docstring'''

											A__   				=					(1, self.in_channels, self.sample_size, self.sample_size)
											A__   				=					jnp.zeros(__lowerCAmelCase      ,				dtype=jnp.floataa)
											A__   				=					jnp.ones((1,)      ,				dtype=jnp.intaa)
											A__   				=					jnp.zeros((1, 1, self.cross_attention_dim)      ,				dtype=jnp.floataa)

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

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


				def  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->Optional[int]:





											'''simple docstring'''

											A__   				=					self.block_out_channels
											A__   				=					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.
											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(__lowerCAmelCase      ,				dtype=self.dtype)

											A__   				=					self.only_cross_attention
											if isinstance(__lowerCAmelCase      ,				__lowerCAmelCase):
																		A__   				=					(only_cross_attention,) * len(self.down_block_types)

											if isinstance(__lowerCAmelCase      ,				__lowerCAmelCase):
																		A__   				=					(num_attention_heads,) * len(self.down_block_types)

											# down
											A__   				=					[]
											A__   				=					block_out_channels[0]
											for i, down_block_type in enumerate(self.down_block_types):
																		A__   				=					output_channel
																		A__   				=					block_out_channels[i]
																		A__   				=					i == len(__lowerCAmelCase) - 1

																		if down_block_type == "CrossAttnDownBlock2D":
																									A__   				=					FlaxCrossAttnDownBlockaD(
																									    in_channels=__lowerCAmelCase      ,				out_channels=__lowerCAmelCase      ,				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:
																									A__   				=					FlaxDownBlockaD(
																									    in_channels=__lowerCAmelCase      ,				out_channels=__lowerCAmelCase      ,				dropout=self.dropout      ,				num_layers=self.layers_per_block      ,				add_downsample=not is_final_block      ,				dtype=self.dtype      ,				)

																		down_blocks.append(__lowerCAmelCase)
											A__   				=					down_blocks

											# mid
											A__   				=					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
											A__   				=					[]
											A__   				=					list(reversed(__lowerCAmelCase))
											A__   				=					list(reversed(__lowerCAmelCase))
											A__   				=					list(reversed(__lowerCAmelCase))
											A__   				=					reversed_block_out_channels[0]
											for i, up_block_type in enumerate(self.up_block_types):
																		A__   				=					output_channel
																		A__   				=					reversed_block_out_channels[i]
																		A__   				=					reversed_block_out_channels[min(i + 1      ,				len(__lowerCAmelCase) - 1)]

																		A__   				=					i == len(__lowerCAmelCase) - 1

																		if up_block_type == "CrossAttnUpBlock2D":
																									A__   				=					FlaxCrossAttnUpBlockaD(
																									    in_channels=__lowerCAmelCase      ,				out_channels=__lowerCAmelCase      ,				prev_output_channel=__lowerCAmelCase      ,				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:
																									A__   				=					FlaxUpBlockaD(
																									    in_channels=__lowerCAmelCase      ,				out_channels=__lowerCAmelCase      ,				prev_output_channel=__lowerCAmelCase      ,				num_layers=self.layers_per_block + 1      ,				add_upsample=not is_final_block      ,				dropout=self.dropout      ,				dtype=self.dtype      ,				)

																		up_blocks.append(__lowerCAmelCase)
																		A__   				=					output_channel
											A__   				=					up_blocks

											# out
											A__   				=					nn.GroupNorm(num_groups=32      ,				epsilon=1e-5)
											A__   				=					nn.Conv(
											    self.out_channels      ,				kernel_size=(3, 3)      ,				strides=(1, 1)      ,				padding=((1, 1), (1, 1))      ,				dtype=self.dtype      ,				)





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





											'''simple docstring'''

											if not isinstance(__lowerCAmelCase      ,				jnp.ndarray):
																		A__   				=					jnp.array([timesteps]      ,				dtype=jnp.intaa)
											elif isinstance(__lowerCAmelCase      ,				jnp.ndarray) and len(timesteps.shape) == 0:
																		A__   				=					timesteps.astype(dtype=jnp.floataa)
																		A__   				=					jnp.expand_dims(__lowerCAmelCase      ,				0)

											A__   				=					self.time_proj(__lowerCAmelCase)
											A__   				=					self.time_embedding(__lowerCAmelCase)

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

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

											if down_block_additional_residuals is not None:
																		A__   				=					()

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

																		A__   				=					new_down_block_res_samples

											# 4. mid
											A__   				=					self.mid_block(__lowerCAmelCase      ,				__lowerCAmelCase      ,				__lowerCAmelCase      ,				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:
																		A__   				=					down_block_res_samples[-(self.layers_per_block + 1) :]
																		A__   				=					down_block_res_samples[: -(self.layers_per_block + 1)]
																		if isinstance(__lowerCAmelCase      ,				__lowerCAmelCase):
																									A__   				=					up_block(
																									    __lowerCAmelCase      ,				temb=__lowerCAmelCase      ,				encoder_hidden_states=__lowerCAmelCase      ,				res_hidden_states_tuple=__lowerCAmelCase      ,				deterministic=not train      ,				)
																		else:
																									A__   				=					up_block(__lowerCAmelCase      ,				temb=__lowerCAmelCase      ,				res_hidden_states_tuple=__lowerCAmelCase      ,				deterministic=not train)

        # 6. post-process
											A__   				=					self.conv_norm_out(__lowerCAmelCase)
											A__   				=					nn.silu(__lowerCAmelCase)
											A__   				=					self.conv_out(__lowerCAmelCase)
											A__   				=					jnp.transpose(__lowerCAmelCase      ,				(0, 3, 1, 2))

											if not return_dict:
																		return (sample,)

											return FlaxUNetaDConditionOutput(sample=__lowerCAmelCase)








import random
import unittest

import torch

from diffusers import IFInpaintingPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin


@skip_mps
class        UpperCamelCase_					(				UpperCAmelCase__	,      UpperCAmelCase__	,      unittest.TestCase  ):

    '''simple docstring'''






    UpperCAmelCase__													=   IFInpaintingPipeline
    UpperCAmelCase__													=   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'''width''', '''height'''}
    UpperCAmelCase__													=   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
    UpperCAmelCase__													=   PipelineTesterMixin.required_optional_params - {'''latents'''}


    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->Dict:





           '''simple docstring'''

           return self._get_dummy_components()


    def  SCREAMING_SNAKE_CASE							(    self   :			int      ,				UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			int=0)		->int:





           '''simple docstring'''

           if str(UpperCAmelCase__).startswith('''mps'''):
                  A__   				=					torch.manual_seed(UpperCAmelCase__)
           else:
                  A__   				=					torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)

           A__   				=					floats_tensor((1, 3, 32, 32)      ,				rng=random.Random(UpperCAmelCase__)).to(UpperCAmelCase__)
           A__   				=					floats_tensor((1, 3, 32, 32)      ,				rng=random.Random(UpperCAmelCase__)).to(UpperCAmelCase__)

           A__   				=					{
               '''prompt''': '''A painting of a squirrel eating a burger''',
               '''image''': image,
               '''mask_image''': mask_image,
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''output_type''': '''numpy''',
           }

           return inputs


    @unittest.skipIf(
        torch_device != '''cuda''' or not is_xformers_available()      ,				reason='''XFormers attention is only available with CUDA and `xformers` installed'''      ,				)
    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->List[Any]:





           '''simple docstring'''

           self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3)


    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->int:





           '''simple docstring'''

           self._test_save_load_optional_components()


    @unittest.skipIf(torch_device != '''cuda'''      ,				reason='''float16 requires CUDA''')
    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any])		->List[str]:





           '''simple docstring'''

           super().test_save_load_floataa(expected_max_diff=1e-1)


    def  SCREAMING_SNAKE_CASE							(    self   :			int)		->List[Any]:





           '''simple docstring'''

           self._test_attention_slicing_forward_pass(expected_max_diff=1e-2)


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





           '''simple docstring'''

           self._test_save_load_local()





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





           '''simple docstring'''

           self._test_inference_batch_single_identical(
               expected_max_diff=1e-2      ,				)


