

"""simple docstring"""




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







@skip_mps
class   __lowercase      (     __lowerCamelCase						,      __lowerCamelCase						,      unittest.TestCase       ):
  snake_case_												=							StableDiffusionPanoramaPipeline
  snake_case_												=							TEXT_TO_IMAGE_PARAMS
  snake_case_												=							TEXT_TO_IMAGE_BATCH_PARAMS
  snake_case_												=							TEXT_TO_IMAGE_IMAGE_PARAMS
  snake_case_												=							TEXT_TO_IMAGE_IMAGE_PARAMS

  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   torch.manual_seed(0     )
   UpperCAmelCase__       :	Tuple        =       UNetaDConditionModel(
       block_out_channels=(32, 64)       ,layers_per_block=1       ,sample_size=32       ,in_channels=4       ,out_channels=4       ,down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")       ,cross_attention_dim=32       ,)
   UpperCAmelCase__       :	int        =       DDIMScheduler()
   torch.manual_seed(0     )
   UpperCAmelCase__       :	str        =       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       ,)
   torch.manual_seed(0     )
   UpperCAmelCase__       :	Optional[int]        =       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=1_000       ,)
   UpperCAmelCase__       :	Optional[Any]        =       CLIPTextModel(A     )
   UpperCAmelCase__       :	List[Any]        =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

   UpperCAmelCase__       :	List[str]        =       {
       """unet""": unet,
       """scheduler""": scheduler,
       """vae""": vae,
       """text_encoder""": text_encoder,
       """tokenizer""": tokenizer,
       """safety_checker""": None,
       """feature_extractor""": None,
   }
   return components

  def  __lowercase							(	self		:      Union[str, Any]       ,A		:      Dict       ,A		:      Tuple=0     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       torch.manual_seed(A     )
   UpperCAmelCase__       :	Union[str, Any]        =       {
       """prompt""": """a photo of the dolomites""",
       """generator""": generator,
       # Setting height and width to None to prevent OOMs on CPU.
       """height""": None,
       """width""": None,
       """num_inference_steps""": 1,
       """guidance_scale""": 6.0,
       """output_type""": """numpy""",
   }
   return inputs

  def  __lowercase							(	self		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase__       :	str        =       self.get_dummy_components()
   UpperCAmelCase__       :	Any        =       StableDiffusionPanoramaPipeline(**A     )
   UpperCAmelCase__       :	Dict        =       sd_pipe.to(A     )
   sd_pipe.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	Union[str, Any]        =       self.get_dummy_inputs(A     )
   UpperCAmelCase__       :	Dict        =       sd_pipe(**A     ).images
   UpperCAmelCase__       :	str        =       image[0, -3:, -3:, -1]
   assert image.shape == (1, 64, 64, 3)

   UpperCAmelCase__       :	Union[str, Any]        =       np.array([0.6_1_8_6, 0.5_3_7_4, 0.4_9_1_5, 0.4_1_3_5, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_7, 0.4_7_5_7]     )

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

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   super().test_inference_batch_consistent(batch_sizes=[1, 2]     )

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   super().test_inference_batch_single_identical(batch_size=2       ,expected_max_diff=3.25e-3     )

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase__       :	str        =       self.get_dummy_components()
   UpperCAmelCase__       :	Union[str, Any]        =       StableDiffusionPanoramaPipeline(**A     )
   UpperCAmelCase__       :	Optional[Any]        =       sd_pipe.to(A     )
   sd_pipe.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	List[Any]        =       self.get_dummy_inputs(A     )
   UpperCAmelCase__       :	Optional[int]        =       """french fries"""
   UpperCAmelCase__       :	List[str]        =       sd_pipe(**A       ,negative_prompt=A     )
   UpperCAmelCase__       :	Optional[Any]        =       output.images
   UpperCAmelCase__       :	Any        =       image[0, -3:, -3:, -1]

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

   UpperCAmelCase__       :	str        =       np.array([0.6_1_8_7, 0.5_3_7_5, 0.4_9_1_5, 0.4_1_3_6, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_6, 0.4_7_5_7]     )

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

  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase__       :	str        =       self.get_dummy_components()
   UpperCAmelCase__       :	Dict        =       StableDiffusionPanoramaPipeline(**A     )
   UpperCAmelCase__       :	List[str]        =       sd_pipe.to(A     )
   sd_pipe.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	Any        =       self.get_dummy_inputs(A     )
   UpperCAmelCase__       :	List[str]        =       sd_pipe(**A       ,view_batch_size=2     )
   UpperCAmelCase__       :	Dict        =       output.images
   UpperCAmelCase__       :	int        =       image[0, -3:, -3:, -1]

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

   UpperCAmelCase__       :	Optional[Any]        =       np.array([0.6_1_8_7, 0.5_3_7_5, 0.4_9_1_5, 0.4_1_3_6, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_6, 0.4_7_5_7]     )

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

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase__       :	Optional[int]        =       self.get_dummy_components()
   UpperCAmelCase__       :	List[str]        =       EulerAncestralDiscreteScheduler(
       beta_start=0.0_0_0_8_5       ,beta_end=0.0_1_2       ,beta_schedule="""scaled_linear"""     )
   UpperCAmelCase__       :	Optional[Any]        =       StableDiffusionPanoramaPipeline(**A     )
   UpperCAmelCase__       :	List[str]        =       sd_pipe.to(A     )
   sd_pipe.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	Optional[int]        =       self.get_dummy_inputs(A     )
   UpperCAmelCase__       :	Tuple        =       sd_pipe(**A     ).images
   UpperCAmelCase__       :	Dict        =       image[0, -3:, -3:, -1]

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

   UpperCAmelCase__       :	Any        =       np.array([0.4_0_2_4, 0.6_5_1_0, 0.4_9_0_1, 0.5_3_7_8, 0.5_8_1_3, 0.5_6_2_2, 0.4_7_9_5, 0.4_4_6_7, 0.4_9_5_2]     )

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




  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
   UpperCAmelCase__       :	List[Any]        =       self.get_dummy_components()
   UpperCAmelCase__       :	Union[str, Any]        =       PNDMScheduler(
       beta_start=0.0_0_0_8_5       ,beta_end=0.0_1_2       ,beta_schedule="""scaled_linear"""       ,skip_prk_steps=A     )
   UpperCAmelCase__       :	List[str]        =       StableDiffusionPanoramaPipeline(**A     )
   UpperCAmelCase__       :	int        =       sd_pipe.to(A     )
   sd_pipe.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	List[str]        =       self.get_dummy_inputs(A     )
   UpperCAmelCase__       :	int        =       sd_pipe(**A     ).images
   UpperCAmelCase__       :	Tuple        =       image[0, -3:, -3:, -1]

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

   UpperCAmelCase__       :	str        =       np.array([0.6_3_9_1, 0.6_2_9_1, 0.4_8_6_1, 0.5_1_3_4, 0.5_5_5_2, 0.4_5_7_8, 0.5_0_3_2, 0.5_0_2_3, 0.4_5_3_9]     )

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







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

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()

  def  __lowercase							(	self		:      int       ,A		:      Dict=0     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       torch.manual_seed(A     )
   UpperCAmelCase__       :	Tuple        =       {
       """prompt""": """a photo of the dolomites""",
       """generator""": generator,
       """num_inference_steps""": 3,
       """guidance_scale""": 7.5,
       """output_type""": """numpy""",
   }
   return inputs

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Union[str, Any]        =       """stabilityai/stable-diffusion-2-base"""
   UpperCAmelCase__       :	Union[str, Any]        =       DDIMScheduler.from_pretrained(A       ,subfolder="""scheduler"""     )
   UpperCAmelCase__       :	List[str]        =       StableDiffusionPanoramaPipeline.from_pretrained(A       ,scheduler=A       ,safety_checker=A     )
   pipe.to(A     )
   pipe.set_progress_bar_config(disable=A     )
   pipe.enable_attention_slicing()

   UpperCAmelCase__       :	Union[str, Any]        =       self.get_inputs()
   UpperCAmelCase__       :	Tuple        =       pipe(**A     ).images
   UpperCAmelCase__       :	Tuple        =       image[0, -3:, -3:, -1].flatten()

   assert image.shape == (1, 512, 2_048, 3)

   UpperCAmelCase__       :	List[Any]        =       np.array(
       [
           0.3_6_9_6_8_3_9_2,
           0.2_7_0_2_5_3_7_2,
           0.3_2_4_4_6_7_6_6,
           0.2_8_3_7_9_3_8_7,
           0.3_6_3_6_3_2_7_4,
           0.3_0_7_3_3_3_4_7,
           0.2_7_1_0_0_0_2_7,
           0.2_7_0_5_4_1_2_5,
           0.2_5_5_3_6_0_9_6,
       ]     )

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

  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       StableDiffusionPanoramaPipeline.from_pretrained(
       """stabilityai/stable-diffusion-2-base"""       ,safety_checker=A     )
   UpperCAmelCase__       :	Optional[int]        =       LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
   pipe.to(A     )
   pipe.set_progress_bar_config(disable=A     )
   pipe.enable_attention_slicing()

   UpperCAmelCase__       :	List[str]        =       self.get_inputs()
   UpperCAmelCase__       :	Union[str, Any]        =       pipe(**A     ).images
   UpperCAmelCase__       :	List[str]        =       image[0, -3:, -3:, -1].flatten()

   assert image.shape == (1, 512, 2_048, 3)

   UpperCAmelCase__       :	Any        =       np.array(
       [
           [
               0.0,
               0.0,
               0.0,
               0.0,
               0.0,
               0.0,
               0.0,
               0.0,
               0.0,
           ]
       ]     )

   assert np.abs(expected_slice - image_slice     ).max() < 1e-3

  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       0

   def callback_fn(A		:      int       ,A		:      int       ,A		:      torch.FloatTensor     ) -> None:
    UpperCAmelCase__       :	Union[str, Any]        =       True
    nonlocal number_of_steps
    number_of_steps += 1
    if step == 1:
     UpperCAmelCase__       :	Tuple        =       latents.detach().cpu().numpy()
     assert latents.shape == (1, 4, 64, 256)
     UpperCAmelCase__       :	Union[str, Any]        =       latents[0, -3:, -3:, -1]

     UpperCAmelCase__       :	Any        =       np.array(
         [
             0.1_8_6_8_1_8_6_9,
             0.3_3_9_0_7_8_1_6,
             0.5_3_6_1_2_7_6,
             0.1_4_4_3_2_8_6_5,
             -0.0_2_8_5_6_6_1_1,
             -0.7_3_9_4_1_1_2_3,
             0.2_3_3_9_7_9_8_7,
             0.4_7_3_2_2_6_8_2,
             -0.3_7_8_2_3_1_6_4,
         ]     )
     assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 5e-2
    elif step == 2:
     UpperCAmelCase__       :	Any        =       latents.detach().cpu().numpy()
     assert latents.shape == (1, 4, 64, 256)
     UpperCAmelCase__       :	Union[str, Any]        =       latents[0, -3:, -3:, -1]

     UpperCAmelCase__       :	int        =       np.array(
         [
             0.1_8_5_3_9_6_4_5,
             0.3_3_9_8_7_2_4_8,
             0.5_3_7_8_5_5_9,
             0.1_4_4_3_7_1_4_2,
             -0.0_2_4_5_5_2_6_1,
             -0.7_3_3_8_3_1_7,
             0.2_3_9_9_0_7_5_5,
             0.4_7_3_5_6_2_7_2,
             -0.3_7_8_6_5_0_5,
         ]     )

     assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 5e-2

   UpperCAmelCase__       :	str        =       False

   UpperCAmelCase__       :	Union[str, Any]        =       """stabilityai/stable-diffusion-2-base"""
   UpperCAmelCase__       :	Optional[int]        =       DDIMScheduler.from_pretrained(A       ,subfolder="""scheduler"""     )
   UpperCAmelCase__       :	int        =       StableDiffusionPanoramaPipeline.from_pretrained(A       ,scheduler=A       ,safety_checker=A     )
   UpperCAmelCase__       :	List[Any]        =       pipe.to(A     )
   pipe.set_progress_bar_config(disable=A     )
   pipe.enable_attention_slicing()

   UpperCAmelCase__       :	List[str]        =       self.get_inputs()
   pipe(**A       ,callback=A       ,callback_steps=1     )
   assert callback_fn.has_been_called
   assert number_of_steps == 3




  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()

   UpperCAmelCase__       :	int        =       """stabilityai/stable-diffusion-2-base"""
   UpperCAmelCase__       :	List[str]        =       DDIMScheduler.from_pretrained(A       ,subfolder="""scheduler"""     )
   UpperCAmelCase__       :	Dict        =       StableDiffusionPanoramaPipeline.from_pretrained(A       ,scheduler=A       ,safety_checker=A     )
   UpperCAmelCase__       :	List[str]        =       pipe.to(A     )
   pipe.set_progress_bar_config(disable=A     )
   pipe.enable_attention_slicing(1     )
   pipe.enable_sequential_cpu_offload()

   UpperCAmelCase__       :	str        =       self.get_inputs()
   UpperCAmelCase__       :	Union[str, Any]        =       pipe(**A     )

   UpperCAmelCase__       :	int        =       torch.cuda.max_memory_allocated()
   # make sure that less than 5.2 GB is allocated
   assert mem_bytes < 5.5 * 10**9









"""simple docstring"""

from pathlib import Path

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


def lowerCAmelCase_   (					UpperCamelCase__						:			np.ndarray					, UpperCamelCase__						:			np.ndarray					, UpperCamelCase__						:			np.ndarray					, UpperCamelCase__						:			int					, UpperCamelCase__						:			int			):






 """simple docstring"""







 __lowercase						=						cva.getAffineTransform(UpperCamelCase__					, UpperCamelCase__			)
 return cva.warpAffine(UpperCamelCase__					, UpperCamelCase__					, (rows, cols)			)


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

       # set different points to rotate image
       UpperCAmelCase__					=np.array([[50, 50], [200, 50], [50, 200]], np.floataa)
       UpperCAmelCase__					=np.array([[10, 100], [200, 50], [100, 250]], np.floataa)
       UpperCAmelCase__					=np.array([[50, 50], [150, 50], [120, 200]], np.floataa)
       UpperCAmelCase__					=np.array([[10, 100], [80, 50], [180, 250]], np.floataa)

       # add all rotated images in a list
       UpperCAmelCase__					=[
           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
       UpperCAmelCase__					=plt.figure(1)
       UpperCAmelCase__					=["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()
