

"""simple docstring"""





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, 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  _lowerCAmelCase  (     lowercase					,unittest.TestCase       ):







 """simple docstring"""







 __UpperCAmelCase     : List[str]       		=     KandinskyVaaControlnetImgaImgPipeline
 __UpperCAmelCase     : Tuple       		=     ["image_embeds", "negative_image_embeds", "image", "hint"]
 __UpperCAmelCase     : Dict       		=     ["image_embeds", "negative_image_embeds", "image", "hint"]
 __UpperCAmelCase     : str       		=     [
     "generator",
     "height",
     "width",
     "strength",
     "guidance_scale",
     "num_inference_steps",
     "return_dict",
     "guidance_scale",
     "num_images_per_prompt",
     "output_type",
     "return_dict",
 ]
 __UpperCAmelCase     : Dict       		=     False


 @property
 def  _lowercase			(      self							:      List[str]     ):
    return 3_2


 @property
 def  _lowercase			(      self							:      Dict     ):
    return 3_2


 @property
 def  _lowercase			(      self							:      Dict     ):
    return self.time_input_dim


 @property
 def  _lowercase			(      self							:      Any     ):
    return self.time_input_dim * 4


 @property
 def  _lowercase			(      self							:      Any     ):
    return 1_0_0


 @property
 def  _lowercase			(      self							:      List[Any]     ):
    torch.manual_seed(0     )

    __lowercase								=						{
        "in_channels": 8,
        # Out channels is double in channels because predicts mean and variance
        "out_channels": 8,
        "addition_embed_type": "image_hint",
        "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,
    }

    __lowercase								=						UNetaDConditionModel(**UpperCAmelCase__     )
    return model


 @property
 def  _lowercase			(      self							:      Optional[int]     ):
    return {
        "block_out_channels": [3_2, 3_2, 6_4, 6_4],
        "down_block_types": [
            "DownEncoderBlock2D",
            "DownEncoderBlock2D",
            "DownEncoderBlock2D",
            "AttnDownEncoderBlock2D",
        ],
        "in_channels": 3,
        "latent_channels": 4,
        "layers_per_block": 1,
        "norm_num_groups": 8,
        "norm_type": "spatial",
        "num_vq_embeddings": 1_2,
        "out_channels": 3,
        "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
        "vq_embed_dim": 4,
    }


 @property
 def  _lowercase			(      self							:      int     ):
    torch.manual_seed(0     )
    __lowercase								=						VQModel(**self.dummy_movq_kwargs     )
    return model


 def  _lowercase			(      self							:      List[Any]     ):
    __lowercase								=						self.dummy_unet
    __lowercase								=						self.dummy_movq

    __lowercase								=						{
        "num_train_timesteps": 1_0_0_0,
        "beta_schedule": "linear",
        "beta_start": 0.00_085,
        "beta_end": 0.012,
        "clip_sample": False,
        "set_alpha_to_one": False,
        "steps_offset": 0,
        "prediction_type": "epsilon",
        "thresholding": False,
    }

    __lowercase								=						DDIMScheduler(**UpperCAmelCase__     )

    __lowercase								=						{
        "unet": unet,
        "scheduler": scheduler,
        "movq": movq,
    }

    return components


 def  _lowercase			(      self							:      Any, UpperCAmelCase__							:      Any, UpperCAmelCase__							:      List[str]=0     ):
    __lowercase								=						floats_tensor((1, self.text_embedder_hidden_size), rng=random.Random(UpperCAmelCase__     )     ).to(UpperCAmelCase__     )
    __lowercase								=						floats_tensor((1, self.text_embedder_hidden_size), rng=random.Random(seed + 1     )     ).to(
        UpperCAmelCase__     )
    # create init_image
    __lowercase								=						floats_tensor((1, 3, 6_4, 6_4), rng=random.Random(UpperCAmelCase__     )     ).to(UpperCAmelCase__     )
    __lowercase								=						image.cpu().permute(0, 2, 3, 1     )[0]
    __lowercase								=						Image.fromarray(np.uinta(UpperCAmelCase__     )     ).convert("RGB"     ).resize((2_5_6, 2_5_6)     )
    # create hint
    __lowercase								=						floats_tensor((1, 3, 6_4, 6_4), rng=random.Random(UpperCAmelCase__     )     ).to(UpperCAmelCase__     )

    if str(UpperCAmelCase__     ).startswith("mps"     ):
       __lowercase								=						torch.manual_seed(UpperCAmelCase__     )
    else:
       __lowercase								=						torch.Generator(device=UpperCAmelCase__     ).manual_seed(UpperCAmelCase__     )
    __lowercase								=						{
        "image": init_image,
        "image_embeds": image_embeds,
        "negative_image_embeds": negative_image_embeds,
        "hint": hint,
        "generator": generator,
        "height": 6_4,
        "width": 6_4,
        "num_inference_steps": 1_0,
        "guidance_scale": 7.0,
        "strength": 0.2,
        "output_type": "np",
    }
    return inputs



 def  _lowercase			(      self							:      str     ):
    __lowercase								=						"cpu"

    __lowercase								=						self.get_dummy_components()

    __lowercase								=						self.pipeline_class(**UpperCAmelCase__     )
    __lowercase								=						pipe.to(UpperCAmelCase__     )

    pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

    __lowercase								=						pipe(**self.get_dummy_inputs(UpperCAmelCase__     )     )
    __lowercase								=						output.images

    __lowercase								=						pipe(
        **self.get_dummy_inputs(UpperCAmelCase__     ), return_dict=UpperCAmelCase__, )[0]

    __lowercase								=						image[0, -3:, -3:, -1]
    __lowercase								=						image_from_tuple[0, -3:, -3:, -1]

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

    __lowercase								=						np.array(
        [0.54_985_034, 0.55_509_365, 0.52_561_504, 0.5_570_494, 0.5_593_818, 0.5_263_979, 0.50_285_643, 0.5_069_846, 0.51_196_736]     )
    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  _lowerCAmelCase  (     unittest.TestCase       ):







 """simple docstring"""







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



 def  _lowercase			(      self							:      List[str]     ):
    __lowercase								=						load_numpy(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        "/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy"     )

    __lowercase								=						load_image(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"     )
    __lowercase								=						init_image.resize((5_1_2, 5_1_2)     )

    __lowercase								=						load_image(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        "/kandinskyv22/hint_image_cat.png"     )
    __lowercase								=						torch.from_numpy(np.array(UpperCAmelCase__     )     ).float() / 255.0
    __lowercase								=						hint.permute(2, 0, 1     ).unsqueeze(0     )

    __lowercase								=						"A robot, 4k photo"

    __lowercase								=						KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
        "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.floataa     )
    pipe_prior.to(UpperCAmelCase__     )

    __lowercase								=						KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
        "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.floataa     )
    __lowercase								=						pipeline.to(UpperCAmelCase__     )

    pipeline.set_progress_bar_config(disable=UpperCAmelCase__     )

    __lowercase								=						torch.Generator(device="cpu"     ).manual_seed(0     )

    __lowercase		,__lowercase								=						pipe_prior(
        UpperCAmelCase__, image=UpperCAmelCase__, strength=0.85, generator=UpperCAmelCase__, negative_prompt="", ).to_tuple()

    __lowercase								=						pipeline(
        image=UpperCAmelCase__, image_embeds=UpperCAmelCase__, negative_image_embeds=UpperCAmelCase__, hint=UpperCAmelCase__, generator=UpperCAmelCase__, num_inference_steps=1_0_0, height=5_1_2, width=5_1_2, strength=0.5, output_type="np", )

    __lowercase								=						output.images[0]

    assert image.shape == (5_1_2, 5_1_2, 3)

    assert_mean_pixel_difference(UpperCAmelCase__, UpperCAmelCase__     )







"""simple docstring"""





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


_a							=							logging.get_logger(__name__)




class  _lowerCAmelCase  (     lowercase       ):







 """simple docstring"""







 __UpperCAmelCase     : Optional[Any]       		=     ["pixel_values"]


 def __init__(      self							:      int, UpperCAmelCase__							:      bool = True, UpperCAmelCase__							:      Optional[Dict[str, int]] = None, UpperCAmelCase__							:      PILImageResampling = PILImageResampling.BILINEAR, UpperCAmelCase__							:      bool = True, UpperCAmelCase__							:      Dict[str, int] = None, UpperCAmelCase__							:      bool = True, UpperCAmelCase__							:      Union[int, float] = 1 / 2_5_5, UpperCAmelCase__							:      bool = True, UpperCAmelCase__							:      Optional[Union[float, List[float]]] = None, UpperCAmelCase__							:      Optional[Union[float, List[float]]] = None, **UpperCAmelCase__							:      str, ):
    super().__init__(**UpperCAmelCase__     )
    __lowercase								=						size if size is not None else {"shortest_edge": 2_5_6}
    __lowercase								=						get_size_dict(UpperCAmelCase__, default_to_square=UpperCAmelCase__     )
    __lowercase								=						crop_size if crop_size is not None else {"height": 2_2_4, "width": 2_2_4}
    __lowercase								=						get_size_dict(UpperCAmelCase__     )
    __lowercase								=						do_resize
    __lowercase								=						size
    __lowercase								=						resample
    __lowercase								=						do_center_crop
    __lowercase								=						crop_size
    __lowercase								=						do_rescale
    __lowercase								=						rescale_factor
    __lowercase								=						do_normalize
    __lowercase								=						image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
    __lowercase								=						image_std if image_std is not None else IMAGENET_STANDARD_STD


 def  _lowercase			(      self							:      int, UpperCAmelCase__							:      np.ndarray, UpperCAmelCase__							:      Dict[str, int], UpperCAmelCase__							:      PILImageResampling = PILImageResampling.BICUBIC, UpperCAmelCase__							:      Optional[Union[str, ChannelDimension]] = None, **UpperCAmelCase__							:      Dict, ):
    __lowercase								=						get_size_dict(UpperCAmelCase__, default_to_square=UpperCAmelCase__     )
    if "shortest_edge" not in size:
       raise ValueError(F"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""     )
    __lowercase								=						get_resize_output_image_size(UpperCAmelCase__, size=size["shortest_edge"], default_to_square=UpperCAmelCase__     )
    return resize(UpperCAmelCase__, size=UpperCAmelCase__, resample=UpperCAmelCase__, data_format=UpperCAmelCase__, **UpperCAmelCase__     )


 def  _lowercase			(      self							:      Dict, UpperCAmelCase__							:      np.ndarray, UpperCAmelCase__							:      Dict[str, int], UpperCAmelCase__							:      Optional[Union[str, ChannelDimension]] = None, **UpperCAmelCase__							:      Dict, ):
    __lowercase								=						get_size_dict(UpperCAmelCase__     )
    return center_crop(UpperCAmelCase__, size=(size["height"], size["width"]), data_format=UpperCAmelCase__, **UpperCAmelCase__     )


 def  _lowercase			(      self							:      Optional[int], UpperCAmelCase__							:      np.ndarray, UpperCAmelCase__							:      float, UpperCAmelCase__							:      Optional[Union[str, ChannelDimension]] = None, **UpperCAmelCase__							:      Union[str, Any]     ):
    return rescale(UpperCAmelCase__, scale=UpperCAmelCase__, data_format=UpperCAmelCase__, **UpperCAmelCase__     )


 def  _lowercase			(      self							:      List[Any], UpperCAmelCase__							:      np.ndarray, UpperCAmelCase__							:      Union[float, List[float]], UpperCAmelCase__							:      Union[float, List[float]], UpperCAmelCase__							:      Optional[Union[str, ChannelDimension]] = None, **UpperCAmelCase__							:      int, ):
    return normalize(UpperCAmelCase__, mean=UpperCAmelCase__, std=UpperCAmelCase__, data_format=UpperCAmelCase__, **UpperCAmelCase__     )



 def  _lowercase			(      self							:      Any, UpperCAmelCase__							:      ImageInput, UpperCAmelCase__							:      Optional[bool] = None, UpperCAmelCase__							:      Dict[str, int] = None, UpperCAmelCase__							:      PILImageResampling = None, UpperCAmelCase__							:      bool = None, UpperCAmelCase__							:      Dict[str, int] = None, UpperCAmelCase__							:      Optional[bool] = None, UpperCAmelCase__							:      Optional[float] = None, UpperCAmelCase__							:      Optional[bool] = None, UpperCAmelCase__							:      Optional[Union[float, List[float]]] = None, UpperCAmelCase__							:      Optional[Union[float, List[float]]] = None, UpperCAmelCase__							:      Optional[Union[str, TensorType]] = None, UpperCAmelCase__							:      Union[str, ChannelDimension] = ChannelDimension.FIRST, **UpperCAmelCase__							:      Optional[int], ):
    __lowercase								=						do_resize if do_resize is not None else self.do_resize
    __lowercase								=						size if size is not None else self.size
    __lowercase								=						get_size_dict(UpperCAmelCase__, default_to_square=UpperCAmelCase__     )
    __lowercase								=						resample if resample is not None else self.resample
    __lowercase								=						do_center_crop if do_center_crop is not None else self.do_center_crop
    __lowercase								=						crop_size if crop_size is not None else self.crop_size
    __lowercase								=						get_size_dict(UpperCAmelCase__     )
    __lowercase								=						do_rescale if do_rescale is not None else self.do_rescale
    __lowercase								=						rescale_factor if rescale_factor is not None else self.rescale_factor
    __lowercase								=						do_normalize if do_normalize is not None else self.do_normalize
    __lowercase								=						image_mean if image_mean is not None else self.image_mean
    __lowercase								=						image_std if image_std is not None else self.image_std

    __lowercase								=						make_list_of_images(UpperCAmelCase__     )

    if not valid_images(UpperCAmelCase__     ):
       raise ValueError(
           "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
           "torch.Tensor, tf.Tensor or jax.ndarray."     )

    if do_resize and size is None:
       raise ValueError("Size must be specified if do_resize is True."     )

    if do_center_crop and crop_size is None:
       raise ValueError("Crop size must be specified if do_center_crop is True."     )

    if do_rescale and rescale_factor is None:
       raise ValueError("Rescale factor must be specified if do_rescale is True."     )

    if do_normalize and (image_mean is None or image_std is None):
       raise ValueError("Image mean and std must be specified if do_normalize is True."     )

    # All transformations expect numpy arrays.
    __lowercase								=						[to_numpy_array(UpperCAmelCase__     ) for image in images]

    if do_resize:
       __lowercase								=						[self.resize(image=UpperCAmelCase__, size=UpperCAmelCase__, resample=UpperCAmelCase__     ) for image in images]

    if do_center_crop:
       __lowercase								=						[self.center_crop(image=UpperCAmelCase__, size=UpperCAmelCase__     ) for image in images]

    if do_rescale:
       __lowercase								=						[self.rescale(image=UpperCAmelCase__, scale=UpperCAmelCase__     ) for image in images]

    if do_normalize:
       __lowercase								=						[self.normalize(image=UpperCAmelCase__, mean=UpperCAmelCase__, std=UpperCAmelCase__     ) for image in images]

    __lowercase								=						[to_channel_dimension_format(UpperCAmelCase__, UpperCAmelCase__     ) for image in images]

    __lowercase								=						{"pixel_values": images}
    return BatchFeature(data=UpperCAmelCase__, tensor_type=UpperCAmelCase__     )




