





"""simple docstring"""


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    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   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase		,	unittest.TestCase ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Optional[Any]						  =       KandinskyVaaImgaImgPipeline
  SCREAMING_SNAKE_CASE__				:Union[str, Any]						  =       ["image_embeds", "negative_image_embeds", "image"]
  SCREAMING_SNAKE_CASE__				:int						  =       [
      "image_embeds",
      "negative_image_embeds",
      "image",
  ]
  SCREAMING_SNAKE_CASE__				:str						  =       [
      "generator",
      "height",
      "width",
      "strength",
      "guidance_scale",
      "num_inference_steps",
      "return_dict",
      "guidance_scale",
      "num_images_per_prompt",
      "output_type",
      "return_dict",
  ]
  SCREAMING_SNAKE_CASE__				:Union[str, Any]						  =       False

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					Optional[int]:
        return 32

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[Any]						)		->					Optional[Any]:
        return 32

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[int]						)		->					Dict:
        return self.time_input_dim

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					Optional[int]:
        return self.time_input_dim * 4

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Tuple						)		->					Any:
        return 100

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					Tuple:
        torch.manual_seed(0						)

        _UpperCamelCase   :						List[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,
        }

        _UpperCamelCase   :						List[str]    =							UNetaDConditionModel(**__a						)
        return model

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					List[str]:
        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  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					str:
        torch.manual_seed(0						)
        _UpperCamelCase   :						Dict    =							VQModel(**self.dummy_movq_kwargs						)
        return model

  def  __SCREAMING_SNAKE_CASE					(		self :							Any						)		->					Tuple:
        _UpperCamelCase   :						Optional[Any]    =							self.dummy_unet
        _UpperCamelCase   :						int    =							self.dummy_movq

        _UpperCamelCase   :						Any    =							{
            "num_train_timesteps": 1000,
            "beta_schedule": "linear",
            "beta_start": 0.0_00_85,
            "beta_end": 0.0_12,
            "clip_sample": False,
            "set_alpha_to_one": False,
            "steps_offset": 0,
            "prediction_type": "epsilon",
            "thresholding": False,
        }

        _UpperCamelCase   :						List[Any]    =							DDIMScheduler(**__a						)

        _UpperCamelCase   :						str    =							{
            "unet": unet,
            "scheduler": scheduler,
            "movq": movq,
        }

        return components

  def  __SCREAMING_SNAKE_CASE					(		self :							int       , __a :							Any       , __a :							List[str]=0						)		->					Tuple:
        _UpperCamelCase   :						Dict    =							floats_tensor((1, self.text_embedder_hidden_size)       , rng=random.Random(__a						)						).to(__a						)
        _UpperCamelCase   :						Dict    =							floats_tensor((1, self.text_embedder_hidden_size)       , rng=random.Random(seed + 1						)						).to(
            __a						)
        # create init_image
        _UpperCamelCase   :						Tuple    =							floats_tensor((1, 3, 64, 64)       , rng=random.Random(__a						)						).to(__a						)
        _UpperCamelCase   :						Any    =							image.cpu().permute(0       , 2       , 3       , 1						)[0]
        _UpperCamelCase   :						Dict    =							Image.fromarray(np.uinta(__a						)						).convert("RGB"						).resize((256, 256)						)

        if str(__a						).startswith("mps"						):
              _UpperCamelCase   :						int    =							torch.manual_seed(__a						)
        else:
              _UpperCamelCase   :						List[str]    =							torch.Generator(device=__a						).manual_seed(__a						)
        _UpperCamelCase   :						Optional[Any]    =							{
            "image": init_image,
            "image_embeds": image_embeds,
            "negative_image_embeds": negative_image_embeds,
            "generator": generator,
            "height": 64,
            "width": 64,
            "num_inference_steps": 10,
            "guidance_scale": 7.0,
            "strength": 0.2,
            "output_type": "np",
        }
        return inputs


  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[int]						)		->					int:
        _UpperCamelCase   :						int    =							"cpu"

        _UpperCamelCase   :						Union[str, Any]    =							self.get_dummy_components()

        _UpperCamelCase   :						Dict    =							self.pipeline_class(**__a						)
        _UpperCamelCase   :						Union[str, Any]    =							pipe.to(__a						)

        pipe.set_progress_bar_config(disable=__a						)

        _UpperCamelCase   :						str    =							pipe(**self.get_dummy_inputs(__a						)						)
        _UpperCamelCase   :						Union[str, Any]    =							output.images

        _UpperCamelCase   :						Any    =							pipe(
            **self.get_dummy_inputs(__a						)       , return_dict=__a       , )[0]

        _UpperCamelCase   :						Optional[int]    =							image[0, -3:, -3:, -1]
        _UpperCamelCase   :						Optional[Any]    =							image_from_tuple[0, -3:, -3:, -1]

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

        _UpperCamelCase   :						Optional[int]    =							np.array(
            [0.6_19_97_78, 0.63_98_44_06, 0.46_14_57_85, 0.62_94_49_84, 0.5_62_22_15, 0.47_30_61_32, 0.47_44_14_56, 0.4_60_76_06, 0.48_71_92_63]						)
        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   __SCREAMING_SNAKE_CASE     (    unittest.TestCase ):







  '''simple docstring'''


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


  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					Dict:
        _UpperCamelCase   :						int    =							load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/kandinskyv22/kandinskyv22_img2img_frog.npy"						)

        _UpperCamelCase   :						Union[str, Any]    =							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"						)
        _UpperCamelCase   :						Union[str, Any]    =							"A red cartoon frog, 4k"

        _UpperCamelCase   :						str    =							KandinskyVaaPriorPipeline.from_pretrained(
            "kandinsky-community/kandinsky-2-2-prior"       , torch_dtype=torch.floataa						)
        pipe_prior.to(__a						)

        _UpperCamelCase   :						List[str]    =							KandinskyVaaImgaImgPipeline.from_pretrained(
            "kandinsky-community/kandinsky-2-2-decoder"       , torch_dtype=torch.floataa						)
        _UpperCamelCase   :						int    =							pipeline.to(__a						)

        pipeline.set_progress_bar_config(disable=__a						)

        _UpperCamelCase   :						List[Any]    =							torch.Generator(device="cpu"						).manual_seed(0						)
        _UpperCamelCase,	_UpperCamelCase   :						Tuple    =							pipe_prior(
            __a       , generator=__a       , num_inference_steps=5       , negative_prompt=""       , ).to_tuple()

        _UpperCamelCase   :						List[Any]    =							pipeline(
            image=__a       , image_embeds=__a       , negative_image_embeds=__a       , generator=__a       , num_inference_steps=100       , height=768       , width=768       , strength=0.2       , output_type="np"       , )

        _UpperCamelCase   :						Optional[Any]    =							output.images[0]

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

        assert_mean_pixel_difference(__a       , __a						)







"""simple docstring"""


from collections.abc import Callable

import numpy as np




def        lowercase__     (       lowercase_		,lowercase_		,lowercase_		,lowercase_		,lowercase_      )       ->		np.ndarray:




      """simple docstring"""

      _UpperCamelCase   :						List[str]    =							int(np.ceil((x_end - xa) / step_size      )      )
      _UpperCamelCase   :						Tuple    =							np.zeros((n + 1,)      )
      _UpperCamelCase   :						List[str]    =							ya
      _UpperCamelCase   :						List[Any]    =							xa

      for k in range(lowercase_      ):
            _UpperCamelCase   :						List[str]    =							y[k] + step_size * ode_func(lowercase_		,y[k]      )
            x += step_size

      return y


if __name__ == "__main__":
  import doctest

  doctest.testmod()
