





import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()







class    __snake_case	(							__lowercase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					KandinskyVaaPipeline
       lowerCAmelCase__						=					[
           '''image_embeds''',
           '''negative_image_embeds''',
       ]
       lowerCAmelCase__						=					['''image_embeds''', '''negative_image_embeds''']
       lowerCAmelCase__						=					[
           '''generator''',
           '''height''',
           '''width''',
           '''latents''',
           '''guidance_scale''',
           '''num_inference_steps''',
           '''return_dict''',
           '''guidance_scale''',
           '''num_images_per_prompt''',
           '''output_type''',
           '''return_dict''',
       ]
       lowerCAmelCase__						=					False




       @property
       def 							UpperCAmelCase__			(  self			:   Dict		):
         return 32




       @property
       def 							UpperCAmelCase__			(  self			:   Dict		):
         return 32




       @property
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         return self.time_input_dim




       @property
       def 							UpperCAmelCase__			(  self			:   Dict		):
         return self.time_input_dim * 4




       @property
       def 							UpperCAmelCase__			(  self			:   Dict		):
         return 100




       @property
       def 							UpperCAmelCase__			(  self			:   Dict		):
         torch.manual_seed(0		)

         __snake_case:						str			  =						{
             """in_channels""": 4,
             # Out channels is double in channels because predicts mean and variance
             """out_channels""": 8,
             """addition_embed_type""": """image""",
             """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
             """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
             """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
             """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
             """layers_per_block""": 1,
             """encoder_hid_dim""": self.text_embedder_hidden_size,
             """encoder_hid_dim_type""": """image_proj""",
             """cross_attention_dim""": self.cross_attention_dim,
             """attention_head_dim""": 4,
             """resnet_time_scale_shift""": """scale_shift""",
             """class_embed_type""": None,
         }

         __snake_case:						str			  =						UNetaDConditionModel(**UpperCAmelCase__		)
         return model




       @property
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         return {
             "block_out_channels": [32, 64],
             "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
             "in_channels": 3,
             "latent_channels": 4,
             "layers_per_block": 1,
             "norm_num_groups": 8,
             "norm_type": "spatial",
             "num_vq_embeddings": 12,
             "out_channels": 3,
             "up_block_types": [
                 "AttnUpDecoderBlock2D",
                 "UpDecoderBlock2D",
             ],
             "vq_embed_dim": 4,
         }




       @property
       def 							UpperCAmelCase__			(  self			:   int		):
         torch.manual_seed(0		)
         __snake_case:						Optional[Any]			  =						VQModel(**self.dummy_movq_kwargs		)
         return model




       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						List[str]			  =						self.dummy_unet
         __snake_case:						Tuple			  =						self.dummy_movq

         __snake_case:						List[str]			  =						DDIMScheduler(
             num_train_timesteps=1_000							,    beta_schedule="""linear"""							,    beta_start=0.0_0085							,    beta_end=0.012							,    clip_sample=UpperCAmelCase__							,    set_alpha_to_one=UpperCAmelCase__							,    steps_offset=1							,    prediction_type="""epsilon"""							,    thresholding=UpperCAmelCase__							,    )

         __snake_case:						Optional[Any]			  =						{
             """unet""": unet,
             """scheduler""": scheduler,
             """movq""": movq,
         }
         return components




       def 							UpperCAmelCase__			(  self			:   Dict							,    A			:   Any							,    A			:   Optional[Any]=0		):
         __snake_case:						Any			  =						floats_tensor((1, self.text_embedder_hidden_size)							,    rng=random.Random(UpperCAmelCase__		)		).to(UpperCAmelCase__		)
         __snake_case:						Optional[Any]			  =						floats_tensor((1, self.text_embedder_hidden_size)							,    rng=random.Random(seed + 1		)		).to(
             UpperCAmelCase__		)
         if str(UpperCAmelCase__		).startswith("""mps"""		):
           __snake_case:						Any			  =						torch.manual_seed(UpperCAmelCase__		)
         else:
           __snake_case:						Union[str, Any]			  =						torch.Generator(device=UpperCAmelCase__		).manual_seed(UpperCAmelCase__		)
         __snake_case:						Optional[int]			  =						{
             """image_embeds""": image_embeds,
             """negative_image_embeds""": negative_image_embeds,
             """generator""": generator,
             """height""": 64,
             """width""": 64,
             """guidance_scale""": 4.0,
             """num_inference_steps""": 2,
             """output_type""": """np""",
         }
         return inputs






       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						Tuple			  =						"""cpu"""

         __snake_case:						int			  =						self.get_dummy_components()

         __snake_case:						Optional[Any]			  =						self.pipeline_class(**UpperCAmelCase__		)
         __snake_case:						str			  =						pipe.to(UpperCAmelCase__		)

         pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

         __snake_case:						Optional[int]			  =						pipe(**self.get_dummy_inputs(UpperCAmelCase__		)		)
         __snake_case:						str			  =						output.images

         __snake_case:						str			  =						pipe(
             **self.get_dummy_inputs(UpperCAmelCase__		)							,    return_dict=UpperCAmelCase__							,    )[0]

         __snake_case:						int			  =						image[0, -3:, -3:, -1]
         __snake_case:						int			  =						image_from_tuple[0, -3:, -3:, -1]

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

         __snake_case:						Dict			  =						np.array(
             [0.623_7976, 1.0, 0.3644_1332, 1.0, 0.7063_9634, 0.2987_7186, 0.8565_2125, 0.521_6843, 0.5445_4046]		)

         assert (
             np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
         ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''

         assert (
             np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2
         ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''







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





       '''simple docstring'''







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






       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						List[Any]			  =						load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
             """/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy"""		)

         __snake_case:						Optional[int]			  =						KandinskyVaaPriorPipeline.from_pretrained(
             """kandinsky-community/kandinsky-2-2-prior"""							,    torch_dtype=torch.floataa		)
         pipe_prior.to(UpperCAmelCase__		)

         __snake_case:						str			  =						KandinskyVaaPipeline.from_pretrained(
             """kandinsky-community/kandinsky-2-2-decoder"""							,    torch_dtype=torch.floataa		)
         __snake_case:						List[Any]			  =						pipeline.to(UpperCAmelCase__		)
         pipeline.set_progress_bar_config(disable=UpperCAmelCase__		)

         __snake_case:						List[str]			  =						"""red cat, 4k photo"""

         __snake_case:						List[Any]			  =						torch.Generator(device="""cuda"""		).manual_seed(0		)
         __snake_case  ,       __snake_case:						int			  =						pipe_prior(
             UpperCAmelCase__							,    generator=UpperCAmelCase__							,    num_inference_steps=5							,    negative_prompt=""""""							,    ).to_tuple()

         __snake_case:						Optional[int]			  =						torch.Generator(device="""cuda"""		).manual_seed(0		)
         __snake_case:						int			  =						pipeline(
             image_embeds=UpperCAmelCase__							,    negative_image_embeds=UpperCAmelCase__							,    generator=UpperCAmelCase__							,    num_inference_steps=100							,    output_type="""np"""							,    )

         __snake_case:						Any			  =						output.images[0]

         assert image.shape == (512, 512, 3)

         assert_mean_pixel_difference(UpperCAmelCase__							,    UpperCAmelCase__		)













from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder







@dataclass
class    __snake_case	(							__lowerCamelCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					42







class    __snake_case	(							__lowerCamelCase       ,      __lowerCamelCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					True




       @register_to_config
       def __init__(  self			:   Union[str, Any]							,    A			:   int = 3							,    A			:   int = 3							,    A			:   Tuple[str] = ("DownEncoderBlock2D",)							,    A			:   Tuple[str] = ("UpDecoderBlock2D",)							,    A			:   Tuple[int] = (64,)							,    A			:   int = 1							,    A			:   str = "silu"							,    A			:   int = 4							,    A			:   int = 32							,    A			:   int = 32							,    A			:   float = 0.1_8215							,    ):
         super().__init__()

         # pass init params to Encoder
         __snake_case:						Any			  =						Encoder(
             in_channels=A							,    out_channels=A							,    down_block_types=A							,    block_out_channels=A							,    layers_per_block=A							,    act_fn=A							,    norm_num_groups=A							,    double_z=A							,    )

         # pass init params to Decoder
         __snake_case:						int			  =						Decoder(
             in_channels=A							,    out_channels=A							,    up_block_types=A							,    block_out_channels=A							,    layers_per_block=A							,    norm_num_groups=A							,    act_fn=A							,    )

         __snake_case:						Dict			  =						nn.Convad(2 * latent_channels							,    2 * latent_channels							,    1		)
         __snake_case:						int			  =						nn.Convad(A							,    A							,    1		)

         __snake_case:						List[str]			  =						False
         __snake_case:						Optional[int]			  =						False

         # only relevant if vae tiling is enabled
         __snake_case:						Any			  =						self.config.sample_size
         __snake_case:						int			  =						(
             self.config.sample_size[0]
             if isinstance(self.config.sample_size							,    (list, tuple)		)
             else self.config.sample_size
         )
         __snake_case:						Union[str, Any]			  =						int(sample_size / (2 ** (len(self.config.block_out_channels		) - 1))		)
         __snake_case:						Optional[int]			  =						0.25




       def 							UpperCAmelCase__			(  self			:   int							,    A			:   List[str]							,    A			:   Optional[Any]=False		):
         if isinstance(A							,    (Encoder, Decoder)		):
           __snake_case:						str			  =						value




       def 							UpperCAmelCase__			(  self			:   str							,    A			:   bool = True		):
         __snake_case:						Union[str, Any]			  =						use_tiling




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         self.enable_tiling(A		)




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						List[str]			  =						True




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case:						List[str]			  =						False




       @property
       # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Any			  =						{}

         def fn_recursive_add_processors(A			:   str							,    A			:   torch.nn.Module							,    A			:   Dict[str, AttentionProcessor]		):
           if hasattr(A							,    """set_processor"""		):
             __snake_case:						List[Any]			  =						module.processor

           for sub_name, child in module.named_children():
             fn_recursive_add_processors(f'''{name}.{sub_name}'''							,    A							,    A		)

           return processors

         for name, module in self.named_children():
           fn_recursive_add_processors(A							,    A							,    A		)

         return processors




       def 							UpperCAmelCase__			(  self			:   Optional[int]							,    A			:   Union[AttentionProcessor, Dict[str, AttentionProcessor]]		):
         __snake_case:						Any			  =						len(self.attn_processors.keys()		)

         if isinstance(A							,    A		) and len(A		) != count:
           raise ValueError(
               f'''A dict of processors was passed, but the number of processors {len(A		)} does not match the'''
               f''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''		)

         def fn_recursive_attn_processor(A			:   str							,    A			:   torch.nn.Module							,    A			:   Optional[Any]		):
           if hasattr(A							,    """set_processor"""		):
             if not isinstance(A							,    A		):
               module.set_processor(A		)
             else:
               module.set_processor(processor.pop(f'''{name}.processor'''		)		)

           for sub_name, child in module.named_children():
             fn_recursive_attn_processor(f'''{name}.{sub_name}'''							,    A							,    A		)

         for name, module in self.named_children():
           fn_recursive_attn_processor(A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         self.set_attn_processor(AttnProcessor()		)




       @apply_forward_hook
       def 							UpperCAmelCase__			(  self			:   Optional[Any]							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
           return self.tiled_encode(A							,    return_dict=A		)

         if self.use_slicing and x.shape[0] > 1:
           __snake_case:						List[Any]			  =						[self.encoder(A		) for x_slice in x.split(1		)]
           __snake_case:						Optional[Any]			  =						torch.cat(A		)
         else:
           __snake_case:						str			  =						self.encoder(A		)

         __snake_case:						Any			  =						self.quant_conv(A		)
         __snake_case:						Tuple			  =						DiagonalGaussianDistribution(A		)

         if not return_dict:
           return (posterior,)

         return AutoencoderKLOutput(latent_dist=A		)




       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
           return self.tiled_decode(A							,    return_dict=A		)

         __snake_case:						Optional[int]			  =						self.post_quant_conv(A		)
         __snake_case:						Union[str, Any]			  =						self.decoder(A		)

         if not return_dict:
           return (dec,)

         return DecoderOutput(sample=A		)




       @apply_forward_hook
       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         if self.use_slicing and z.shape[0] > 1:
           __snake_case:						Union[str, Any]			  =						[self._decode(A		).sample for z_slice in z.split(1		)]
           __snake_case:						List[str]			  =						torch.cat(A		)
         else:
           __snake_case:						int			  =						self._decode(A		).sample

         if not return_dict:
           return (decoded,)

         return DecoderOutput(sample=A		)




       def 							UpperCAmelCase__			(  self			:   Any							,    A			:   Tuple							,    A			:   int							,    A			:   List[Any]		):
         __snake_case:						int			  =						min(a.shape[2]							,    b.shape[2]							,    A		)
         for y in range(A		):
           __snake_case:						Dict			  =						a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
         return b




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   Optional[Any]							,    A			:   List[str]							,    A			:   List[str]		):
         __snake_case:						Dict			  =						min(a.shape[3]							,    b.shape[3]							,    A		)
         for x in range(A		):
           __snake_case:						Tuple			  =						a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
         return b




       def 							UpperCAmelCase__			(  self			:   int							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         __snake_case:						List[str]			  =						int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)		)
         __snake_case:						Dict			  =						int(self.tile_latent_min_size * self.tile_overlap_factor		)
         __snake_case:						Dict			  =						self.tile_latent_min_size - blend_extent

         # Split the image into 512x512 tiles and encode them separately.
         __snake_case:						Optional[int]			  =						[]
         for i in range(0							,    x.shape[2]							,    A		):
           __snake_case:						Optional[int]			  =						[]
           for j in range(0							,    x.shape[3]							,    A		):
             __snake_case:						int			  =						x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
             __snake_case:						Tuple			  =						self.encoder(A		)
             __snake_case:						Dict			  =						self.quant_conv(A		)
             row.append(A		)
           rows.append(A		)
         __snake_case:						Tuple			  =						[]
         for i, row in enumerate(A		):
           __snake_case:						str			  =						[]
           for j, tile in enumerate(A		):
             # blend the above tile and the left tile
             # to the current tile and add the current tile to the result row
             if i > 0:
               __snake_case:						Optional[Any]			  =						self.blend_v(rows[i - 1][j]							,    A							,    A		)
             if j > 0:
               __snake_case:						Tuple			  =						self.blend_h(row[j - 1]							,    A							,    A		)
             result_row.append(tile[:, :, :row_limit, :row_limit]		)
           result_rows.append(torch.cat(A							,    dim=3		)		)

         __snake_case:						Tuple			  =						torch.cat(A							,    dim=2		)
         __snake_case:						Optional[int]			  =						DiagonalGaussianDistribution(A		)

         if not return_dict:
           return (posterior,)

         return AutoencoderKLOutput(latent_dist=A		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         __snake_case:						Optional[Any]			  =						int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)		)
         __snake_case:						str			  =						int(self.tile_sample_min_size * self.tile_overlap_factor		)
         __snake_case:						int			  =						self.tile_sample_min_size - blend_extent

         # Split z into overlapping 64x64 tiles and decode them separately.
         # The tiles have an overlap to avoid seams between tiles.
         __snake_case:						List[Any]			  =						[]
         for i in range(0							,    z.shape[2]							,    A		):
           __snake_case:						Optional[Any]			  =						[]
           for j in range(0							,    z.shape[3]							,    A		):
             __snake_case:						Dict			  =						z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
             __snake_case:						Any			  =						self.post_quant_conv(A		)
             __snake_case:						Optional[Any]			  =						self.decoder(A		)
             row.append(A		)
           rows.append(A		)
         __snake_case:						Optional[Any]			  =						[]
         for i, row in enumerate(A		):
           __snake_case:						Optional[Any]			  =						[]
           for j, tile in enumerate(A		):
             # blend the above tile and the left tile
             # to the current tile and add the current tile to the result row
             if i > 0:
               __snake_case:						Tuple			  =						self.blend_v(rows[i - 1][j]							,    A							,    A		)
             if j > 0:
               __snake_case:						List[str]			  =						self.blend_h(row[j - 1]							,    A							,    A		)
             result_row.append(tile[:, :, :row_limit, :row_limit]		)
           result_rows.append(torch.cat(A							,    dim=3		)		)

         __snake_case:						Dict			  =						torch.cat(A							,    dim=2		)
         if not return_dict:
           return (dec,)

         return DecoderOutput(sample=A		)






       def 							UpperCAmelCase__			(  self			:   List[Any]							,    A			:   torch.FloatTensor							,    A			:   bool = False							,    A			:   bool = True							,    A			:   Optional[torch.Generator] = None							,    ):
         __snake_case:						Optional[Any]			  =						sample
         __snake_case:						Union[str, Any]			  =						self.encode(A		).latent_dist
         if sample_posterior:
           __snake_case:						Optional[Any]			  =						posterior.sample(generator=A		)
         else:
           __snake_case:						Dict			  =						posterior.mode()
         __snake_case:						Any			  =						self.decode(A		).sample

         if not return_dict:
           return (dec,)

         return DecoderOutput(sample=A		)






