






'''simple docstring'''







import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class      __SCREAMING_SNAKE_CASE					(__A       ,					__A       ,					__A       ,					unittest.TestCase					):


  """simple docstring"""


  _a     :   List[str]			 =    StableDiffusionImageVariationPipeline
  _a     :   Optional[int]			 =    IMAGE_VARIATION_PARAMS
  _a     :   Any			 =    IMAGE_VARIATION_BATCH_PARAMS
  _a     :   List[str]			 =    frozenset([]					)
  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  _a     :   Optional[int]			 =    frozenset([]					)







  def       _a					( self  ):






         """simple docstring"""


         torch.manual_seed(0  )
         a_	    = 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')      ,  cross_attention_dim=32      ,  )
         a_	    = PNDMScheduler(skip_prk_steps=UpperCamelCase__  )
         torch.manual_seed(0  )
         a_	    = 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  )
         a_	    = CLIPVisionConfig(
             hidden_size=32      ,  projection_dim=32      ,  intermediate_size=37      ,  layer_norm_eps=1e-05      ,  num_attention_heads=4      ,  num_hidden_layers=5      ,  image_size=32      ,  patch_size=4      ,  )
         a_	    = CLIPVisionModelWithProjection(UpperCamelCase__  )
         a_	    = CLIPImageProcessor(crop_size=32      ,  size=32  )

         a_	    = {
             'unet': unet,
             'scheduler': scheduler,
             'vae': vae,
             'image_encoder': image_encoder,
             'feature_extractor': feature_extractor,
             'safety_checker': None,
         }
         return components







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__=0  ):






         """simple docstring"""


         a_	    = floats_tensor((1, 3, 32, 32)      ,  rng=random.Random(UpperCamelCase__  )  )
         a_	    = image.cpu().permute(0      ,  2      ,  3      ,  1  )[0]
         a_	    = Image.fromarray(np.uinta(UpperCamelCase__  )  ).convert('RGB'  ).resize((32, 32)  )
         if str(UpperCamelCase__  ).startswith('mps'  ):
                a_	    = torch.manual_seed(UpperCamelCase__  )
         else:
                a_	    = torch.Generator(device=UpperCamelCase__  ).manual_seed(UpperCamelCase__  )
         a_	    = {
             'image': image,
             'generator': generator,
             'num_inference_steps': 2,
             'guidance_scale': 6.0,
             'output_type': 'numpy',
         }
         return inputs







  def       _a					( self  ):






         """simple docstring"""


         a_	    = 'cpu'  # ensure determinism for the device-dependent torch.Generator
         a_	    = self.get_dummy_components()
         a_	    = StableDiffusionImageVariationPipeline(**UpperCamelCase__  )
         a_	    = sd_pipe.to(UpperCamelCase__  )
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

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

         assert image.shape == (1, 64, 64, 3)
         a_	    = np.array([0.5_239, 0.5_723, 0.4_796, 0.5_049, 0.5_550, 0.4_685, 0.5_329, 0.4_891, 0.4_921]  )

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







  def       _a					( self  ):






         """simple docstring"""


         a_	    = 'cpu'  # ensure determinism for the device-dependent torch.Generator
         a_	    = self.get_dummy_components()
         a_	    = StableDiffusionImageVariationPipeline(**UpperCamelCase__  )
         a_	    = sd_pipe.to(UpperCamelCase__  )
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

         a_	    = self.get_dummy_inputs(UpperCamelCase__  )
         a_	    = 2 * [inputs['image']]
         a_	    = sd_pipe(**UpperCamelCase__  )

         a_	    = output.images

         a_	    = image[-1, -3:, -3:, -1]

         assert image.shape == (2, 64, 64, 3)
         a_	    = np.array([0.6_892, 0.5_637, 0.5_836, 0.5_771, 0.6_254, 0.6_409, 0.5_580, 0.5_569, 0.5_289]  )

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







  def       _a					( self  ):






         """simple docstring"""


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





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


  """simple docstring"""


  def       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__="cpu"      ,  UpperCamelCase__=torch.floataa      ,  UpperCamelCase__=0  ):






         """simple docstring"""


         a_	    = torch.Generator(device=UpperCamelCase__  ).manual_seed(UpperCamelCase__  )
         a_	    = load_image(
             'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
             '/stable_diffusion_imgvar/input_image_vermeer.png'  )
         a_	    = np.random.RandomState(UpperCamelCase__  ).standard_normal((1, 4, 64, 64)  )
         a_	    = torch.from_numpy(UpperCamelCase__  ).to(device=UpperCamelCase__      ,  dtype=UpperCamelCase__  )
         a_	    = {
             'image': init_image,
             'latents': latents,
             'generator': generator,
             'num_inference_steps': 3,
             'guidance_scale': 7.5,
             'output_type': 'numpy',
         }
         return inputs







  def       _a					( self  ):






         """simple docstring"""


         a_	    = StableDiffusionImageVariationPipeline.from_pretrained(
             'lambdalabs/sd-image-variations-diffusers'      ,  safety_checker=UpperCamelCase__  )
         a_	    = sd_pipe.to(UpperCamelCase__  )
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

         a_	    = self.get_inputs(UpperCamelCase__  )
         a_	    = sd_pipe(**UpperCamelCase__  ).images
         a_	    = image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 512, 512, 3)
         a_	    = np.array([0.84_491, 0.90_789, 0.75_708, 0.78_734, 0.83_485, 0.70_099, 0.66_938, 0.68_727, 0.61_379]  )
         assert np.abs(image_slice - expected_slice  ).max() < 6e-3







  def       _a					( self  ):






         """simple docstring"""


         a_	    = 0

         def callback_fn(UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__  ) -> None:
                a_	    = True
                nonlocal number_of_steps
                number_of_steps += 1
                if step == 1:
                       a_	    = latents.detach().cpu().numpy()
                       assert latents.shape == (1, 4, 64, 64)
                       a_	    = latents[0, -3:, -3:, -1]
                       a_	    = np.array(
                           [-0.1_621, 0.2_837, -0.7_979, -0.1_221, -1.3_057, 0.7_681, -2.1_191, 0.0_464, 1.6_309]  )

                       assert np.abs(latents_slice.flatten() - expected_slice  ).max() < 5e-2
                elif step == 2:
                       a_	    = latents.detach().cpu().numpy()
                       assert latents.shape == (1, 4, 64, 64)
                       a_	    = latents[0, -3:, -3:, -1]
                       a_	    = np.array([0.6_299, 1.7_500, 1.1_992, -2.1_582, -1.8_994, 0.7_334, -0.7_090, 1.0_137, 1.5_273]  )

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

         a_	    = False

         a_	    = StableDiffusionImageVariationPipeline.from_pretrained(
             'fusing/sd-image-variations-diffusers'      ,  safety_checker=UpperCamelCase__      ,  torch_dtype=torch.floataa      ,  )
         pipe.to(UpperCamelCase__  )
         pipe.set_progress_bar_config(disable=UpperCamelCase__  )
         pipe.enable_attention_slicing()

         a_	    = self.get_inputs(UpperCamelCase__      ,  dtype=torch.floataa  )
         pipe(**UpperCamelCase__      ,  callback=UpperCamelCase__      ,  callback_steps=1  )
         assert callback_fn.has_been_called
         assert number_of_steps == inputs["num_inference_steps"]







  def       _a					( self  ):






         """simple docstring"""


         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         a_	    = 'fusing/sd-image-variations-diffusers'
         a_	    = StableDiffusionImageVariationPipeline.from_pretrained(
             UpperCamelCase__      ,  safety_checker=UpperCamelCase__      ,  torch_dtype=torch.floataa  )
         a_	    = pipe.to(UpperCamelCase__  )
         pipe.set_progress_bar_config(disable=UpperCamelCase__  )
         pipe.enable_attention_slicing(1  )
         pipe.enable_sequential_cpu_offload()

         a_	    = self.get_inputs(UpperCamelCase__      ,  dtype=torch.floataa  )
         a_	    = pipe(**UpperCamelCase__  )

         a_	    = torch.cuda.max_memory_allocated()
         # make sure that less than 2.6 GB is allocated
         assert mem_bytes < 2.6 * 10**9





@nightly
@require_torch_gpu
class      __SCREAMING_SNAKE_CASE					(unittest.TestCase					):


  """simple docstring"""


  def       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__="cpu"      ,  UpperCamelCase__=torch.floataa      ,  UpperCamelCase__=0  ):






         """simple docstring"""


         a_	    = torch.Generator(device=UpperCamelCase__  ).manual_seed(UpperCamelCase__  )
         a_	    = load_image(
             'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
             '/stable_diffusion_imgvar/input_image_vermeer.png'  )
         a_	    = np.random.RandomState(UpperCamelCase__  ).standard_normal((1, 4, 64, 64)  )
         a_	    = torch.from_numpy(UpperCamelCase__  ).to(device=UpperCamelCase__      ,  dtype=UpperCamelCase__  )
         a_	    = {
             'image': init_image,
             'latents': latents,
             'generator': generator,
             'num_inference_steps': 50,
             'guidance_scale': 7.5,
             'output_type': 'numpy',
         }
         return inputs







  def       _a					( self  ):






         """simple docstring"""


         a_	    = StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'  )
         sd_pipe.to(UpperCamelCase__  )
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

         a_	    = self.get_inputs(UpperCamelCase__  )
         a_	    = sd_pipe(**UpperCamelCase__  ).images[0]

         a_	    = load_numpy(
             'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
             '/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'  )
         a_	    = np.abs(expected_image - image  ).max()
         assert max_diff < 1e-3







  def       _a					( self  ):






         """simple docstring"""


         a_	    = StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'  )
         a_	    = DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config  )
         sd_pipe.to(UpperCamelCase__  )
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

         a_	    = self.get_inputs(UpperCamelCase__  )
         a_	    = 25
         a_	    = sd_pipe(**UpperCamelCase__  ).images[0]

         a_	    = load_numpy(
             'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
             '/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'  )
         a_	    = np.abs(expected_image - image  ).max()
         assert max_diff < 1e-3












'''simple docstring'''







import re


def 			__UpperCamelCase				(			lowercase_					:     int      ):
       """simple docstring"""




       def docstring_decorator(lowercase_					:     int      ):
              a_	    = fn.__doc__
              a_	    = func_doc.split('\n'      )
              a_	    = 0
              while i < len(lowercase_      ) and re.search(R'^\s*Examples?:\s*$'      ,					lines[i]      ) is None:
                     i += 1
              if i < len(lowercase_      ):
                     a_	    = example_docstring
                     a_	    = '\n'.join(lowercase_      )
              else:
                     raise ValueError(
                         F'The function {fn} should have an empty \'Examples:\' in its docstring as placeholder, '
                         F'current docstring is:\n{func_doc}'      )
              a_	    = func_doc
              return fn

       return docstring_decorator




