


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 				snake_case__(							lowercase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  jnp.ndarray
						SCREAMING_SNAKE_CASE__      :  jnp.ndarray







class 				snake_case__(							nn.Module							):







						'''simple docstring'''




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




						def 						lowercase_ (				self							)					->    Optional[int]:
										lowerCAmelCase_ :      int    				=				nn.Conv(
										    self.block_out_channels[0]   ,						kernel_size=(3, 3)   ,						padding=((1, 1), (1, 1))   ,						dtype=self.dtype   ,						)

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

										lowerCAmelCase_ :      List[Any]    				=				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   ,						__lowercase							)					->    Union[str, Any]:
										lowerCAmelCase_ :      Optional[int]    				=				self.conv_in(__lowercase							)
										lowerCAmelCase_ :      Any    				=				nn.silu(__lowercase							)

										for block in self.blocks:
														lowerCAmelCase_ :      int    				=				block(__lowercase							)
														lowerCAmelCase_ :      List[str]    				=				nn.silu(__lowercase							)

										lowerCAmelCase_ :      int    				=				self.conv_out(__lowercase							)

										return embedding







@flax_register_to_config
class 				snake_case__(							nn.Module,						lowercase__,						lowercase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  int					=							32
						SCREAMING_SNAKE_CASE__      :  int					=							4
						SCREAMING_SNAKE_CASE__      :  Tuple[str]					=							(
						    "CrossAttnDownBlock2D",
						    "CrossAttnDownBlock2D",
						    "CrossAttnDownBlock2D",
						    "DownBlock2D",
						)
						SCREAMING_SNAKE_CASE__      :  Union[bool, Tuple[bool]]					=							False
						SCREAMING_SNAKE_CASE__      :  Tuple[int]					=							(320, 640, 1_280, 1_280)
						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					=							1_280
						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__      :  str					=							"rgb"
						SCREAMING_SNAKE_CASE__      :  Tuple[int]					=							(16, 32, 96, 256)




						def 						lowercase_ (				self   ,						__lowercase							)					->    List[str]:
										# init input tensors
										lowerCAmelCase_ :      List[str]    				=				(1, self.in_channels, self.sample_size, self.sample_size)
										lowerCAmelCase_ :      Optional[int]    				=				jnp.zeros(__lowercase   ,						dtype=jnp.floataa							)
										lowerCAmelCase_ :      List[Any]    				=				jnp.ones((1,)   ,						dtype=jnp.intaa							)
										lowerCAmelCase_ :      Optional[int]    				=				jnp.zeros((1, 1, self.cross_attention_dim)   ,						dtype=jnp.floataa							)
										lowerCAmelCase_ :      Dict    				=				(1, 3, self.sample_size * 8, self.sample_size * 8)
										lowerCAmelCase_ :      Optional[Any]    				=				jnp.zeros(__lowercase   ,						dtype=jnp.floataa							)

										lowerCAmelCase_				,  lowerCAmelCase_ :      Any    				=				jax.random.split(__lowercase							)
										lowerCAmelCase_ :      Tuple    				=				{'''params''': params_rng, '''dropout''': dropout_rng}

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




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      List[str]    				=				self.block_out_channels
										lowerCAmelCase_ :      List[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.
										lowerCAmelCase_ :      List[Any]    				=				self.num_attention_heads or self.attention_head_dim

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

										# time
										lowerCAmelCase_ :      Dict    				=				FlaxTimesteps(
										    block_out_channels[0]   ,						flip_sin_to_cos=self.flip_sin_to_cos   ,						freq_shift=self.config.freq_shift							)
										lowerCAmelCase_ :      Union[str, Any]    				=				FlaxTimestepEmbedding(__lowercase   ,						dtype=self.dtype							)

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

										lowerCAmelCase_ :      Dict    				=				self.only_cross_attention
										if isinstance(__lowercase   ,						__lowercase							):
														lowerCAmelCase_ :      int    				=				(only_cross_attention,) * len(self.down_block_types							)

										if isinstance(__lowercase   ,						__lowercase							):
														lowerCAmelCase_ :      Optional[Any]    				=				(num_attention_heads,) * len(self.down_block_types							)

										# down
										lowerCAmelCase_ :      str    				=				[]
										lowerCAmelCase_ :      Any    				=				[]

										lowerCAmelCase_ :      Tuple    				=				block_out_channels[0]

										lowerCAmelCase_ :      Optional[int]    				=				nn.Conv(
										    __lowercase   ,						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(__lowercase							)

										for i, down_block_type in enumerate(self.down_block_types							):
														lowerCAmelCase_ :      Optional[int]    				=				output_channel
														lowerCAmelCase_ :      str    				=				block_out_channels[i]
														lowerCAmelCase_ :      List[str]    				=				i == len(__lowercase							) - 1

														if down_block_type == "CrossAttnDownBlock2D":
																		lowerCAmelCase_ :      List[str]    				=				FlaxCrossAttnDownBlockaD(
																		    in_channels=__lowercase   ,						out_channels=__lowercase   ,						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:
																		lowerCAmelCase_ :      Union[str, Any]    				=				FlaxDownBlockaD(
																		    in_channels=__lowercase   ,						out_channels=__lowercase   ,						dropout=self.dropout   ,						num_layers=self.layers_per_block   ,						add_downsample=not is_final_block   ,						dtype=self.dtype   ,						)

														down_blocks.append(__lowercase							)

														for _ in range(self.layers_per_block							):
																		lowerCAmelCase_ :      Optional[int]    				=				nn.Conv(
																		    __lowercase   ,						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(__lowercase							)

														if not is_final_block:
																		lowerCAmelCase_ :      int    				=				nn.Conv(
																		    __lowercase   ,						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(__lowercase							)

										lowerCAmelCase_ :      List[Any]    				=				down_blocks
										lowerCAmelCase_ :      Tuple    				=				controlnet_down_blocks

										# mid
										lowerCAmelCase_ :      Tuple    				=				block_out_channels[-1]
										lowerCAmelCase_ :      str    				=				FlaxUNetMidBlockaDCrossAttn(
										    in_channels=__lowercase   ,						dropout=self.dropout   ,						num_attention_heads=num_attention_heads[-1]   ,						use_linear_projection=self.use_linear_projection   ,						dtype=self.dtype   ,						)

										lowerCAmelCase_ :      Union[str, Any]    				=				nn.Conv(
										    __lowercase   ,						kernel_size=(1, 1)   ,						padding='''VALID'''   ,						kernel_init=nn.initializers.zeros_init()   ,						bias_init=nn.initializers.zeros_init()   ,						dtype=self.dtype   ,						)




						def __call__(				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase = 1.0   ,						__lowercase = True   ,						__lowercase = False   ,						)					->    Tuple:
										lowerCAmelCase_ :      Any    				=				self.controlnet_conditioning_channel_order
										if channel_order == "bgr":
														lowerCAmelCase_ :      Optional[int]    				=				jnp.flip(__lowercase   ,						axis=1							)

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

										lowerCAmelCase_ :      str    				=				self.time_proj(__lowercase							)
										lowerCAmelCase_ :      Optional[Any]    				=				self.time_embedding(__lowercase							)

										# 2. pre-process
										lowerCAmelCase_ :      Dict    				=				jnp.transpose(__lowercase   ,						(0, 2, 3, 1)							)
										lowerCAmelCase_ :      str    				=				self.conv_in(__lowercase							)

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

										# 3. down
										lowerCAmelCase_ :      Optional[int]    				=				(sample,)
										for down_block in self.down_blocks:
														if isinstance(__lowercase   ,						__lowercase							):
																		lowerCAmelCase_				,  lowerCAmelCase_ :      str    				=				down_block(__lowercase   ,						__lowercase   ,						__lowercase   ,						deterministic=not train							)
														else:
																		lowerCAmelCase_				,  lowerCAmelCase_ :      Tuple    				=				down_block(__lowercase   ,						__lowercase   ,						deterministic=not train							)
														down_block_res_samples += res_samples

										# 4. mid
										lowerCAmelCase_ :      int    				=				self.mid_block(__lowercase   ,						__lowercase   ,						__lowercase   ,						deterministic=not train							)

										# 5. contronet blocks
										lowerCAmelCase_ :      int    				=				()
										for down_block_res_sample, controlnet_block in zip(__lowercase   ,						self.controlnet_down_blocks							):
														lowerCAmelCase_ :      Optional[Any]    				=				controlnet_block(__lowercase							)
														controlnet_down_block_res_samples += (down_block_res_sample,)

										lowerCAmelCase_ :      Union[str, Any]    				=				controlnet_down_block_res_samples

										lowerCAmelCase_ :      Optional[Any]    				=				self.controlnet_mid_block(__lowercase							)

										# 6. scaling
										lowerCAmelCase_ :      Any    				=				[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=__lowercase   ,						mid_block_res_sample=__lowercase							)



import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_UpperCAmelCase			:  Any     =False







class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						@property
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Any:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Optional[Any]:
										return 3_2




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Any    				=				VQModel(
										    block_out_channels=[3_2, 6_4]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,						latent_channels=3   ,						num_vq_embeddings=self.num_embed   ,						vq_embed_dim=3   ,						)
										return model




						@property
						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)
										return tokenizer




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				CLIPTextConfig(
										    bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=self.text_embedder_hidden_size   ,						intermediate_size=3_7   ,						layer_norm_eps=1e-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1_0_0_0   ,						)
										return CLIPTextModel(__lowercase							)




						@property
						def 						lowercase_ (				self							)					->    List[str]:
										torch.manual_seed(0							)

										lowerCAmelCase_ :      Union[str, Any]    				=				1_2
										lowerCAmelCase_ :      int    				=				1_2

										lowerCAmelCase_ :      Union[str, Any]    				=				{
										    '''attention_bias''': True,
										    '''cross_attention_dim''': 3_2,
										    '''attention_head_dim''': height * width,
										    '''num_attention_heads''': 1,
										    '''num_vector_embeds''': self.num_embed,
										    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
										    '''norm_num_groups''': 3_2,
										    '''sample_size''': width,
										    '''activation_fn''': '''geglu-approximate''',
										}

										lowerCAmelCase_ :      List[str]    				=				TransformeraDModel(**__lowercase							)
										return model




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      Any    				=				self.dummy_vqvae
										lowerCAmelCase_ :      str    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      int    				=				self.dummy_transformer
										lowerCAmelCase_ :      List[str]    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      Union[str, Any]    				=				LearnedClassifierFreeSamplingEmbeddings(learnable=__lowercase							)

										lowerCAmelCase_ :      Dict    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      int    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      int    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Tuple    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      Union[str, Any]    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      List[Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 2_4, 2_4, 3)

										lowerCAmelCase_ :      Optional[int]    				=				np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]							)

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




						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      Optional[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      str    				=				self.dummy_vqvae
										lowerCAmelCase_ :      Dict    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      List[Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_transformer
										lowerCAmelCase_ :      Tuple    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      str    				=				LearnedClassifierFreeSamplingEmbeddings(
										    learnable=__lowercase   ,						hidden_size=self.text_embedder_hidden_size   ,						length=tokenizer.model_max_length							)

										lowerCAmelCase_ :      List[str]    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      List[str]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      str    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 2_4, 2_4, 3)

										lowerCAmelCase_ :      Union[str, Any]    				=				np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() < 2.0
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2







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







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      Tuple    				=				load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''							)

										lowerCAmelCase_ :      str    				=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''							)
										lowerCAmelCase_ :      List[Any]    				=				pipeline.to(__lowercase							)
										pipeline.set_progress_bar_config(disable=__lowercase							)

										# requires GPU generator for gumbel softmax
										# don't use GPU generator in tests though
										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Optional[int]    				=				pipeline(
										    '''teddy bear playing in the pool'''   ,						num_images_per_prompt=1   ,						generator=__lowercase   ,						output_type='''np'''   ,						)

										lowerCAmelCase_ :      Union[str, Any]    				=				output.images[0]

										assert image.shape == (2_5_6, 2_5_6, 3)
										assert np.abs(expected_image - image							).max() < 2.0