





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class 				a			(      __lowercase   ,unittest.TestCase			):
				SCREAMING_SNAKE_CASE__     :   List[str]		      =     LDMTextToImagePipeline
				SCREAMING_SNAKE_CASE__     :   int		      =     TEXT_TO_IMAGE_PARAMS - {
				    '''negative_prompt''',
				    '''negative_prompt_embeds''',
				    '''cross_attention_kwargs''',
				    '''prompt_embeds''',
				}
				SCREAMING_SNAKE_CASE__     :   int		      =     PipelineTesterMixin.required_optional_params - {
				    '''num_images_per_prompt''',
				    '''callback''',
				    '''callback_steps''',
				}
				SCREAMING_SNAKE_CASE__     :   Optional[Any]		      =     TEXT_TO_IMAGE_BATCH_PARAMS
				SCREAMING_SNAKE_CASE__     :   int		      =     False






				def 			snake_case_    (  self					):

										"""simple docstring"""




										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				int           =   UNetaDConditionModel(
										    block_out_channels=(32, 64)			, layers_per_block=2			, sample_size=32			, in_channels=4			, out_channels=4			, down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')			, up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')			, cross_attention_dim=32			, )
										__SCREAMING_SNAKE_CASE:				Any           =   DDIMScheduler(
										    beta_start=0.00085			, beta_end=0.012			, beta_schedule='''scaled_linear'''			, clip_sample=_lowerCAmelCase			, set_alpha_to_one=_lowerCAmelCase			, )
										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				Tuple           =   AutoencoderKL(
										    block_out_channels=(32, 64)			, in_channels=3			, out_channels=3			, down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D''')			, up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D''')			, latent_channels=4			, )
										torch.manual_seed(0					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   CLIPTextConfig(
										    bos_token_id=0			, eos_token_id=2			, hidden_size=32			, intermediate_size=37			, layer_norm_eps=1e-05			, num_attention_heads=4			, num_hidden_layers=5			, pad_token_id=1			, vocab_size=1000			, )
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   CLIPTextModel(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Tuple           =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										    '''vqvae''': vae,
										    '''bert''': text_encoder,
										    '''tokenizer''': tokenizer,
										}
										return components






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase=0					):

										"""simple docstring"""




										if str(_lowerCAmelCase					).startswith('''mps'''					):
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.manual_seed(_lowerCAmelCase					)
										else:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.Generator(device=_lowerCAmelCase					).manual_seed(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Dict           =   {
										    '''prompt''': '''A painting of a squirrel eating a burger''',
										    '''generator''': generator,
										    '''num_inference_steps''': 2,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										}
										return inputs






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator

										__SCREAMING_SNAKE_CASE:				str           =   self.get_dummy_components()
										__SCREAMING_SNAKE_CASE:				Any           =   LDMTextToImagePipeline(**_lowerCAmelCase					)
										pipe.to(_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				List[Any]           =   self.get_dummy_inputs(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				List[Any]           =   pipe(**_lowerCAmelCase					).images
										__SCREAMING_SNAKE_CASE:				List[Any]           =   image[0, -3:, -3:, -1]

										assert image.shape == (1, 16, 16, 3)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   np.array([0.6101, 0.6156, 0.5622, 0.4895, 0.6661, 0.3804, 0.5748, 0.6136, 0.5014]					)

										assert np.abs(image_slice.flatten() - expected_slice					).max() < 1e-3





@slow
@require_torch_gpu
class 				a			(      unittest.TestCase			):






				def 			snake_case_    (  self					):

										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase=torch.floataa			, _lowerCAmelCase=0					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   torch.manual_seed(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   np.random.RandomState(_lowerCAmelCase					).standard_normal((1, 4, 32, 32)					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.from_numpy(_lowerCAmelCase					).to(device=_lowerCAmelCase			, dtype=_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {
										    '''prompt''': '''A painting of a squirrel eating a burger''',
										    '''latents''': latents,
										    '''generator''': generator,
										    '''num_inference_steps''': 3,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										}
										return inputs






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256'''					).to(_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				int           =   self.get_inputs(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   pipe(**_lowerCAmelCase					).images
										__SCREAMING_SNAKE_CASE:				int           =   image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 256, 256, 3)
										__SCREAMING_SNAKE_CASE:				Any           =   np.array([0.51825, 0.52850, 0.52543, 0.54258, 0.52304, 0.52569, 0.54363, 0.55276, 0.56878]					)
										__SCREAMING_SNAKE_CASE:				int           =   np.abs(expected_slice - image_slice					).max()
										assert max_diff < 1e-3





@nightly
@require_torch_gpu
class 				a			(      unittest.TestCase			):






				def 			snake_case_    (  self					):

										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase=torch.floataa			, _lowerCAmelCase=0					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   torch.manual_seed(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				List[Any]           =   np.random.RandomState(_lowerCAmelCase					).standard_normal((1, 4, 32, 32)					)
										__SCREAMING_SNAKE_CASE:				int           =   torch.from_numpy(_lowerCAmelCase					).to(device=_lowerCAmelCase			, dtype=_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   {
										    '''prompt''': '''A painting of a squirrel eating a burger''',
										    '''latents''': latents,
										    '''generator''': generator,
										    '''num_inference_steps''': 50,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										}
										return inputs






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256'''					).to(_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				List[Any]           =   self.get_inputs(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Dict           =   pipe(**_lowerCAmelCase					).images[0]

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   load_numpy(
										    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy'''					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   np.abs(expected_image - image					).max()
										assert max_diff < 1e-3










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 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   jnp.ndarray





@flax_register_to_config
class 				a			(      nn.Module   ,__lowercase   ,__lowercase			):
				SCREAMING_SNAKE_CASE__     :   int		      =     32
				SCREAMING_SNAKE_CASE__     :   int		      =     4
				SCREAMING_SNAKE_CASE__     :   int		      =     4
				SCREAMING_SNAKE_CASE__     :   Tuple[str]		      =     (
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				SCREAMING_SNAKE_CASE__     :   Tuple[str]		      =     ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				SCREAMING_SNAKE_CASE__     :   Union[bool, Tuple[bool]]		      =     False
				SCREAMING_SNAKE_CASE__     :   Tuple[int]		      =     (320, 640, 1280, 1280)
				SCREAMING_SNAKE_CASE__     :   int		      =     2
				SCREAMING_SNAKE_CASE__     :   Union[int, Tuple[int]]		      =     8
				SCREAMING_SNAKE_CASE__     :   Optional[Union[int, Tuple[int]]]		      =     None
				SCREAMING_SNAKE_CASE__     :   int		      =     1280
				SCREAMING_SNAKE_CASE__     :   float		      =     0.0
				SCREAMING_SNAKE_CASE__     :   bool		      =     False
				SCREAMING_SNAKE_CASE__     :   jnp.dtype		      =     jnp.floataa
				SCREAMING_SNAKE_CASE__     :   bool		      =     True
				SCREAMING_SNAKE_CASE__     :   int		      =     0
				SCREAMING_SNAKE_CASE__     :   bool		      =     False






				def 			snake_case_    (  self			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   (1, self.in_channels, self.sample_size, self.sample_size)
										__SCREAMING_SNAKE_CASE:				Tuple           =   jnp.zeros(_lowerCAmelCase			, dtype=jnp.floataa					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   jnp.ones((1,)			, dtype=jnp.intaa					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   jnp.zeros((1, 1, self.cross_attention_dim)			, dtype=jnp.floataa					)

										__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Optional[int]           =   jax.random.split(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''params''': params_rng, '''dropout''': dropout_rng}

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   self.block_out_channels
										__SCREAMING_SNAKE_CASE:				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.
										__SCREAMING_SNAKE_CASE:				Any           =   self.num_attention_heads or self.attention_head_dim

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

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

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.only_cross_attention
										if isinstance(_lowerCAmelCase			, _lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   (only_cross_attention,) * len(self.down_block_types					)

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

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

																if down_block_type == "CrossAttnDownBlock2D":
																						__SCREAMING_SNAKE_CASE:				str           =   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:
																						__SCREAMING_SNAKE_CASE:				Tuple           =   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					)
										__SCREAMING_SNAKE_CASE:				str           =   down_blocks

										# mid
										__SCREAMING_SNAKE_CASE:				Union[str, 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
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   []
										__SCREAMING_SNAKE_CASE:				Tuple           =   list(reversed(_lowerCAmelCase					)					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   list(reversed(_lowerCAmelCase					)					)
										__SCREAMING_SNAKE_CASE:				str           =   list(reversed(_lowerCAmelCase					)					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   reversed_block_out_channels[0]
										for i, up_block_type in enumerate(self.up_block_types					):
																__SCREAMING_SNAKE_CASE:				int           =   output_channel
																__SCREAMING_SNAKE_CASE:				List[str]           =   reversed_block_out_channels[i]
																__SCREAMING_SNAKE_CASE:				List[str]           =   reversed_block_out_channels[min(i + 1			, len(_lowerCAmelCase					) - 1					)]

																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   i == len(_lowerCAmelCase					) - 1

																if up_block_type == "CrossAttnUpBlock2D":
																						__SCREAMING_SNAKE_CASE:				Optional[int]           =   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:
																						__SCREAMING_SNAKE_CASE:				int           =   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					)
																__SCREAMING_SNAKE_CASE:				Any           =   output_channel
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   up_blocks

										# out
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   nn.GroupNorm(num_groups=32			, epsilon=1e-5					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   nn.Conv(
										    self.out_channels			, kernel_size=(3, 3)			, strides=(1, 1)			, padding=((1, 1), (1, 1))			, dtype=self.dtype			, )






				def __call__(  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase=None			, _lowerCAmelCase=None			, _lowerCAmelCase = True			, _lowerCAmelCase = False			, ):

										"""simple docstring"""




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

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

										# 2. pre-process
										__SCREAMING_SNAKE_CASE:				int           =   jnp.transpose(_lowerCAmelCase			, (0, 2, 3, 1)					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   self.conv_in(_lowerCAmelCase					)

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

										if down_block_additional_residuals is not None:
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   ()

																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,)

																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   new_down_block_res_samples

										# 4. mid
										__SCREAMING_SNAKE_CASE:				Dict           =   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:
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   down_block_res_samples[-(self.layers_per_block + 1) :]
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   down_block_res_samples[: -(self.layers_per_block + 1)]
																if isinstance(_lowerCAmelCase			, _lowerCAmelCase					):
																						__SCREAMING_SNAKE_CASE:				Any           =   up_block(
																						    _lowerCAmelCase			, temb=_lowerCAmelCase			, encoder_hidden_states=_lowerCAmelCase			, res_hidden_states_tuple=_lowerCAmelCase			, deterministic=not train			, )
																else:
																						__SCREAMING_SNAKE_CASE:				List[str]           =   up_block(_lowerCAmelCase			, temb=_lowerCAmelCase			, res_hidden_states_tuple=_lowerCAmelCase			, deterministic=not train					)

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

										if not return_dict:
																return (sample,)

										return FlaxUNetaDConditionOutput(sample=_lowerCAmelCase					)



