






'''simple docstring'''





import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class 				lowercase       (							_lowerCamelCase ,       _lowerCamelCase ,       unittest.TestCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							StableDiffusionXLImgaImgPipeline
    lowerCAmelCase__      				=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
    lowerCAmelCase__      				=							PipelineTesterMixin.required_optional_params - {"latents"}
    lowerCAmelCase__      				=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowerCAmelCase__      				=							IMAGE_TO_IMAGE_IMAGE_PARAMS
    lowerCAmelCase__      				=							IMAGE_TO_IMAGE_IMAGE_PARAMS

    def    a__				(					self     :     Dict     )							->							Union[str, Any]:

     '''simple docstring'''




     torch.manual_seed(0     )
     lowerCamelCase__  =     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")				,     attention_head_dim=(2, 4)				,     use_linear_projection=__lowerCamelCase				,     addition_embed_type="text_time"				,     addition_time_embed_dim=8				,     transformer_layers_per_block=(1, 2)				,     projection_class_embeddings_input_dim=80				,     cross_attention_dim=64				,     )
     lowerCamelCase__  =     EulerDiscreteScheduler(
         beta_start=0.0_0_0_8_5				,     beta_end=0.0_1_2				,     steps_offset=1				,     beta_schedule="scaled_linear"				,     timestep_spacing="leading"				,     )
     torch.manual_seed(0     )
     lowerCamelCase__  =     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				,     sample_size=128				,     )
     torch.manual_seed(0     )
     lowerCamelCase__  =     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				,     hidden_act="gelu"				,     projection_dim=32				,     )
     lowerCamelCase__  =     CLIPTextModel(__lowerCamelCase     )
     lowerCamelCase__  =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				,     local_files_only=__lowerCamelCase     )

     lowerCamelCase__  =     CLIPTextModelWithProjection(__lowerCamelCase     )
     lowerCamelCase__  =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				,     local_files_only=__lowerCamelCase     )

     lowerCamelCase__  =     {
         "unet": unet,
         "scheduler": scheduler,
         "vae": vae,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "text_encoder_2": text_encoder_a,
         "tokenizer_2": tokenizer_a,
         # "safety_checker": None,
         # "feature_extractor": None,
     }
     return components

    def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Union[str, Any]=0     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     floats_tensor((1, 3, 32, 32)				,     rng=random.Random(__lowerCamelCase     )     ).to(__lowerCamelCase     )
     lowerCamelCase__  =     image / 2 + 0.5
     if str(__lowerCamelCase     ).startswith("mps"     ):
      lowerCamelCase__  =     torch.manual_seed(__lowerCamelCase     )
     else:
      lowerCamelCase__  =     torch.Generator(device=__lowerCamelCase     ).manual_seed(__lowerCamelCase     )
     lowerCamelCase__  =     {
         "prompt": "A painting of a squirrel eating a burger",
         "image": image,
         "generator": generator,
         "num_inference_steps": 2,
         "guidance_scale": 5.0,
         "output_type": "numpy",
         "strength": 0.7_5,
     }
     return inputs

    def    a__				(					self     :     Union[str, Any]     )							->							Tuple:

     '''simple docstring'''




     lowerCamelCase__  =     "cpu"  # ensure determinism for the device-dependent torch.Generator
     lowerCamelCase__  =     self.get_dummy_components()
     lowerCamelCase__  =     StableDiffusionXLImgaImgPipeline(**__lowerCamelCase     )
     lowerCamelCase__  =     sd_pipe.to(__lowerCamelCase     )
     sd_pipe.set_progress_bar_config(disable=__lowerCamelCase     )

     lowerCamelCase__  =     self.get_dummy_inputs(__lowerCamelCase     )
     lowerCamelCase__  =     sd_pipe(**__lowerCamelCase     ).images
     lowerCamelCase__  =     image[0, -3:, -3:, -1]

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

     lowerCamelCase__  =     np.array([0.4_6_5_6, 0.4_8_4_0, 0.4_4_3_9, 0.6_6_9_8, 0.5_5_7_4, 0.4_5_2_4, 0.5_7_9_9, 0.5_9_4_3, 0.5_1_6_5]     )

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

    def    a__				(					self     :     Union[str, Any]     )							->							Union[str, Any]:

     '''simple docstring'''




     super().test_attention_slicing_forward_pass(expected_max_diff=3E-3     )

    def    a__				(					self     :     Optional[Any]     )							->							List[str]:

     '''simple docstring'''




     super().test_inference_batch_single_identical(expected_max_diff=3E-3     )

    def    a__				(					self     :     List[Any]     )							->							Any:

     '''simple docstring'''




     pass




    def    a__				(					self     :     Optional[int]     )							->							List[str]:

     '''simple docstring'''




     lowerCamelCase__  =     self.get_dummy_components()
     lowerCamelCase__  =     StableDiffusionXLImgaImgPipeline(**__lowerCamelCase     )
     lowerCamelCase__  =     sd_pipe.to(__lowerCamelCase     )
     lowerCamelCase__  =     sd_pipe.to(__lowerCamelCase     )
     sd_pipe.set_progress_bar_config(disable=__lowerCamelCase     )

     # forward without prompt embeds
     lowerCamelCase__  =     self.get_dummy_inputs(__lowerCamelCase     )
     lowerCamelCase__  =     3 * ["this is a negative prompt"]
     lowerCamelCase__  =     negative_prompt
     lowerCamelCase__  =     3 * [inputs["prompt"]]

     lowerCamelCase__  =     sd_pipe(**__lowerCamelCase     )
     lowerCamelCase__  =     output.images[0, -3:, -3:, -1]

     # forward with prompt embeds
     lowerCamelCase__  =     self.get_dummy_inputs(__lowerCamelCase     )
     lowerCamelCase__  =     3 * ["this is a negative prompt"]
     lowerCamelCase__  =     3 * [inputs.pop("prompt"     )]

     (
         (
         lowerCamelCase__
     )  ,	(
         lowerCamelCase__
     )  ,	(
         lowerCamelCase__
     )  ,	(
         lowerCamelCase__
     )  ,	
     )  =     sd_pipe.encode_prompt(__lowerCamelCase				,     negative_prompt=__lowerCamelCase     )

     lowerCamelCase__  =     sd_pipe(
         **__lowerCamelCase				,     prompt_embeds=__lowerCamelCase				,     negative_prompt_embeds=__lowerCamelCase				,     pooled_prompt_embeds=__lowerCamelCase				,     negative_pooled_prompt_embeds=__lowerCamelCase				,     )
     lowerCamelCase__  =     output.images[0, -3:, -3:, -1]

     # make sure that it's equal
     assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()     ).max() < 1E-4
@slow
@require_torch_gpu
class 				lowercase       (							unittest.TestCase      ):


    '''simple docstring'''

    def    a__				(					self     :     Optional[int]     )							->							Any:

     '''simple docstring'''




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

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Optional[Any]="cpu"				,     __lowerCamelCase     :     Optional[int]=torch.floataa				,     __lowerCamelCase     :     Union[str, Any]=0     )							->							Tuple:

     '''simple docstring'''




     lowerCamelCase__  =     torch.Generator(device=__lowerCamelCase     ).manual_seed(__lowerCamelCase     )
     lowerCamelCase__  =     np.random.RandomState(__lowerCamelCase     ).standard_normal((1, 4, 64, 64)     )
     lowerCamelCase__  =     torch.from_numpy(__lowerCamelCase     ).to(device=__lowerCamelCase				,     dtype=__lowerCamelCase     )
     lowerCamelCase__  =     {
         "prompt": "a photograph of an astronaut riding a horse",
         "latents": latents,
         "generator": generator,
         "num_inference_steps": 3,
         "guidance_scale": 7.5,
         "output_type": "numpy",
     }
     return inputs




    def    a__				(					self     :     Optional[int]     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"     )
     pipe.to(__lowerCamelCase     )
     pipe.set_progress_bar_config(disable=__lowerCamelCase     )

     lowerCamelCase__  =     self.get_inputs(__lowerCamelCase     )
     lowerCamelCase__  =     pipe(**__lowerCamelCase     ).images
     lowerCamelCase__  =     image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 512, 512, 3)
     lowerCamelCase__  =     np.array([0.4_9_4_9_3, 0.4_7_8_9_6, 0.4_0_7_9_8, 0.5_4_2_1_4, 0.5_3_2_1_2, 0.4_8_2_0_2, 0.4_7_6_5_6, 0.4_6_3_2_9, 0.4_8_5_0_6]     )
     assert np.abs(image_slice - expected_slice     ).max() < 7E-3









'''simple docstring'''





from pathlib import Path

import cva
import numpy as np
from matplotlib import pyplot as plt






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__):
 lowerCamelCase__  =     cva.getAffineTransform(lowercase__				,     lowercase__)
 return cva.warpAffine(lowercase__				,     lowercase__				,     (rows, cols))


if __name__ == "__main__":
 # read original image
 __A							:							Any  							=		cva.imread(
     str(Path(__file__).resolve().parent.parent / """image_data""" / """lena.jpg""")
 )
 # turn image in gray scale value
 __A							:							int  							=		cva.cvtColor(image, cva.COLOR_BGR2GRAY)
 # get image shape
 __A							,			__A							:							Dict  							=		gray_img.shape

 # set different points to rotate image
 __A							:							int  							=		np.array([[50, 50], [2_00, 50], [50, 2_00]], np.floataa)
 __A							:							Any  							=		np.array([[10, 1_00], [2_00, 50], [1_00, 2_50]], np.floataa)
 __A							:							str  							=		np.array([[50, 50], [1_50, 50], [1_20, 2_00]], np.floataa)
 __A							:							Optional[int]  							=		np.array([[10, 1_00], [80, 50], [1_80, 2_50]], np.floataa)

 # add all rotated images in a list
 __A							:							List[Any]  							=		[
     gray_img,
     get_rotation(gray_img, ptsa, ptsa, img_rows, img_cols),
     get_rotation(gray_img, ptsa, ptsa, img_rows, img_cols),
     get_rotation(gray_img, ptsa, ptsa, img_rows, img_cols),
 ]

 # plot different image rotations
 __A							:							int  							=		plt.figure(1)
 __A							:							int  							=		["""Original""", """Rotation 1""", """Rotation 2""", """Rotation 3"""]
 for i, image in enumerate(images):
  plt.subplot(2, 2, i + 1), plt.imshow(image, """gray""")
  plt.title(titles[i])
  plt.axis("""off""")
  plt.subplots_adjust(left=0.0, bottom=0.05, right=1.0, top=0.95)
 plt.show()

