





import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, 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_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class     lowercase__       (	__lowerCamelCase			,				__lowerCamelCase			,				__lowerCamelCase			,				unittest.TestCase				):

   '''simple docstring'''


   a						:							Union[str, Any]   		=		StableDiffusionInstructPixaPixPipeline
   a						:							int   		=		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
   a						:							int   		=		TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
   a						:							Dict   		=		IMAGE_TO_IMAGE_IMAGE_PARAMS
   a						:							Optional[Any]   		=		IMAGE_TO_IMAGE_IMAGE_PARAMS




   def       UpperCamelCase__  (    self							) ->      Optional[Any]:






       """simple docstring"""
       torch.manual_seed(0							)
       UpperCamelCase__							:				Optional[int]          =      UNetaDConditionModel(
           block_out_channels=(32, 64),		layers_per_block=2,		sample_size=32,		in_channels=8,		out_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),		cross_attention_dim=32,		)
       UpperCamelCase__							:				Union[str, Any]          =      PNDMScheduler(skip_prk_steps=__magic_name__							)
       torch.manual_seed(0							)
       UpperCamelCase__							:				Tuple          =      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__							:				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=1000,		)
       UpperCamelCase__							:				str          =      CLIPTextModel(__magic_name__							)
       UpperCamelCase__							:				Optional[Any]          =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

       UpperCamelCase__							:				Optional[Any]          =      {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__=0							) ->      List[str]:






       """simple docstring"""
       UpperCamelCase__							:				Optional[Any]          =      floats_tensor((1, 3, 32, 32),		rng=random.Random(__magic_name__							)							).to(__magic_name__							)
       UpperCamelCase__							:				Dict          =      image.cpu().permute(0,		2,		3,		1							)[0]
       UpperCamelCase__							:				str          =      Image.fromarray(np.uinta(__magic_name__							)							).convert('''RGB'''							)
       if str(__magic_name__							).startswith('''mps'''							):
           UpperCamelCase__							:				Tuple          =      torch.manual_seed(__magic_name__							)
       else:
           UpperCamelCase__							:				List[str]          =      torch.Generator(device=__magic_name__							).manual_seed(__magic_name__							)
       UpperCamelCase__							:				Union[str, Any]          =      {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''image_guidance_scale''': 1,
           '''output_type''': '''numpy''',
       }
       return inputs




   def       UpperCamelCase__  (    self							) ->      int:






       """simple docstring"""
       UpperCamelCase__							:				int          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCamelCase__							:				Optional[Any]          =      self.get_dummy_components()
       UpperCamelCase__							:				int          =      StableDiffusionInstructPixaPixPipeline(**__magic_name__							)
       UpperCamelCase__							:				List[Any]          =      sd_pipe.to(__magic_name__							)
       sd_pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:				Optional[Any]          =      self.get_dummy_inputs(__magic_name__							)
       UpperCamelCase__							:				str          =      sd_pipe(**__magic_name__							).images
       UpperCamelCase__							:				int          =      image[0, -3:, -3:, -1]
       assert image.shape == (1, 32, 32, 3)
       UpperCamelCase__							:				int          =      np.array([0.7526, 0.3750, 0.4547, 0.6117, 0.5866, 0.5016, 0.4327, 0.5642, 0.4815]							)

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




   def       UpperCamelCase__  (    self							) ->      Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:				Optional[Any]          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCamelCase__							:				List[str]          =      self.get_dummy_components()
       UpperCamelCase__							:				Optional[int]          =      StableDiffusionInstructPixaPixPipeline(**__magic_name__							)
       UpperCamelCase__							:				Optional[int]          =      sd_pipe.to(__magic_name__							)
       sd_pipe.set_progress_bar_config(disable=__magic_name__							)

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

       assert image.shape == (1, 32, 32, 3)
       UpperCamelCase__							:				Union[str, Any]          =      np.array([0.7511, 0.3642, 0.4553, 0.6236, 0.5797, 0.5013, 0.4343, 0.5611, 0.4831]							)

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




   def       UpperCamelCase__  (    self							) ->      Any:






       """simple docstring"""
       UpperCamelCase__							:				Dict          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCamelCase__							:				Optional[Any]          =      self.get_dummy_components()
       UpperCamelCase__							:				str          =      StableDiffusionInstructPixaPixPipeline(**__magic_name__							)
       UpperCamelCase__							:				Tuple          =      sd_pipe.to(__magic_name__							)
       sd_pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:				Optional[int]          =      self.get_dummy_inputs(__magic_name__							)
       UpperCamelCase__							:				Tuple          =      [inputs['''prompt''']] * 2

       UpperCamelCase__							:				int          =      np.array(inputs['''image''']							).astype(np.floataa							) / 255.0
       UpperCamelCase__							:				Dict          =      torch.from_numpy(__magic_name__							).unsqueeze(0							).to(__magic_name__							)
       UpperCamelCase__							:				Optional[Any]          =      image / 2 + 0.5
       UpperCamelCase__							:				Optional[int]          =      image.permute(0,		3,		1,		2							)
       UpperCamelCase__							:				List[str]          =      image.repeat(2,		1,		1,		1							)

       UpperCamelCase__							:				Union[str, Any]          =      sd_pipe(**__magic_name__							).images
       UpperCamelCase__							:				int          =      image[-1, -3:, -3:, -1]

       assert image.shape == (2, 32, 32, 3)
       UpperCamelCase__							:				Tuple          =      np.array([0.5812, 0.5748, 0.5222, 0.5908, 0.5695, 0.7174, 0.6804, 0.5523, 0.5579]							)

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




   def       UpperCamelCase__  (    self							) ->      Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:				Union[str, Any]          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCamelCase__							:				Optional[int]          =      self.get_dummy_components()
       UpperCamelCase__							:				int          =      EulerAncestralDiscreteScheduler(
           beta_start=0.0_0085,		beta_end=0.012,		beta_schedule='''scaled_linear'''							)
       UpperCamelCase__							:				str          =      StableDiffusionInstructPixaPixPipeline(**__magic_name__							)
       UpperCamelCase__							:				Dict          =      sd_pipe.to(__magic_name__							)
       sd_pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:				Optional[int]          =      self.get_dummy_inputs(__magic_name__							)
       UpperCamelCase__							:				List[str]          =      sd_pipe(**__magic_name__							).images
       UpperCamelCase__							:				Any          =      image[0, -3:, -3:, -1]

       UpperCamelCase__							:				Any          =      [round(__magic_name__,		4							) for x in image_slice.flatten().tolist()]
       print(''','''.join([str(__magic_name__							) for x in slice]							)							)

       assert image.shape == (1, 32, 32, 3)
       UpperCamelCase__							:				Dict          =      np.array([0.7417, 0.3842, 0.4732, 0.5776, 0.5891, 0.5139, 0.4052, 0.5673, 0.4986]							)

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




   def       UpperCamelCase__  (    self							) ->      Any:






       """simple docstring"""
       super().test_inference_batch_single_identical(expected_max_diff=3E-3							)




   def       UpperCamelCase__  (    self							) ->      Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:				str          =      self.get_dummy_components()
       UpperCamelCase__							:				Dict          =      StableDiffusionInstructPixaPixPipeline(**__magic_name__							)
       UpperCamelCase__							:				int          =      VaeImageProcessor(do_resize=__magic_name__,		do_normalize=__magic_name__							)
       UpperCamelCase__							:				List[Any]          =      pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:				Tuple          =      pipe(**self.get_dummy_inputs_by_type(__magic_name__,		input_image_type='''pt'''							)							)[0]

       UpperCamelCase__							:				Any          =      components['''vae''']
       UpperCamelCase__							:				Tuple          =      self.get_dummy_inputs_by_type(__magic_name__,		input_image_type='''pt'''							)

       for image_param in self.image_latents_params:
           if image_param in inputs.keys():
               UpperCamelCase__							:				List[str]          =      vae.encode(inputs[image_param]							).latent_dist.mode()

       UpperCamelCase__							:				List[Any]          =      pipe(**__magic_name__							)[0]

       UpperCamelCase__							:				int          =      np.abs(out - out_latents_inputs							).max()
       self.assertLess(__magic_name__,		1E-4,		'''passing latents as image input generate different result from passing image'''							)






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

   '''simple docstring'''


   def       UpperCamelCase__  (    self							) ->      int:






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




   def       UpperCamelCase__  (    self,		__magic_name__=0							) ->      Any:






       """simple docstring"""
       UpperCamelCase__							:				Union[str, Any]          =      torch.manual_seed(__magic_name__							)
       UpperCamelCase__							:				str          =      load_image(
           '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'''							)
       UpperCamelCase__							:				Tuple          =      {
           '''prompt''': '''turn him into a cyborg''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''image_guidance_scale''': 1.0,
           '''output_type''': '''numpy''',
       }
       return inputs




   def       UpperCamelCase__  (    self							) ->      int:






       """simple docstring"""
       UpperCamelCase__							:				str          =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix''',		safety_checker=__magic_name__							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)
       pipe.enable_attention_slicing()

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

       assert image.shape == (1, 512, 512, 3)
       UpperCamelCase__							:				str          =      np.array([0.5902, 0.6015, 0.6027, 0.5983, 0.6092, 0.6061, 0.5765, 0.5785, 0.5555]							)

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




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






       """simple docstring"""
       UpperCamelCase__							:				Dict          =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix''',		safety_checker=__magic_name__							)
       UpperCamelCase__							:				int          =      LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)
       pipe.enable_attention_slicing()

       UpperCamelCase__							:				int          =      self.get_inputs()
       UpperCamelCase__							:				str          =      pipe(**__magic_name__							).images
       UpperCamelCase__							:				str          =      image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 512, 3)
       UpperCamelCase__							:				str          =      np.array([0.6578, 0.6817, 0.6972, 0.6761, 0.6856, 0.6916, 0.6428, 0.6516, 0.6301]							)

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




   def       UpperCamelCase__  (    self							) ->      Tuple:






       """simple docstring"""
       UpperCamelCase__							:				Tuple          =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix''',		safety_checker=__magic_name__							)
       UpperCamelCase__							:				Dict          =      DDIMScheduler.from_config(pipe.scheduler.config							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)
       pipe.enable_attention_slicing()

       UpperCamelCase__							:				Tuple          =      self.get_inputs()
       UpperCamelCase__							:				Tuple          =      pipe(**__magic_name__							).images
       UpperCamelCase__							:				List[Any]          =      image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 512, 3)
       UpperCamelCase__							:				Dict          =      np.array([0.3828, 0.3834, 0.3818, 0.3792, 0.3865, 0.3752, 0.3792, 0.3847, 0.3753]							)

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




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






       """simple docstring"""
       UpperCamelCase__							:				List[str]          =      0

       def callback_fn(__magic_name__,		__magic_name__,		__magic_name__							) -> None:
           UpperCamelCase__							:				Optional[Any]          =      True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 1:
               UpperCamelCase__							:				Any          =      latents.detach().cpu().numpy()
               assert latents.shape == (1, 4, 64, 64)
               UpperCamelCase__							:				str          =      latents[0, -3:, -3:, -1]
               UpperCamelCase__							:				str          =      np.array([-0.2463, -0.4644, -0.9756, 1.5176, 1.4414, 0.7866, 0.9897, 0.8521, 0.7983]							)

               assert np.abs(latents_slice.flatten() - expected_slice							).max() < 5E-2
           elif step == 2:
               UpperCamelCase__							:				Tuple          =      latents.detach().cpu().numpy()
               assert latents.shape == (1, 4, 64, 64)
               UpperCamelCase__							:				Tuple          =      latents[0, -3:, -3:, -1]
               UpperCamelCase__							:				int          =      np.array([-0.2644, -0.4626, -0.9653, 1.5176, 1.4551, 0.7686, 0.9805, 0.8452, 0.8115]							)

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

       UpperCamelCase__							:				Any          =      False

       UpperCamelCase__							:				Dict          =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix''',		safety_checker=__magic_name__,		torch_dtype=torch.floataa							)
       UpperCamelCase__							:				Any          =      pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)
       pipe.enable_attention_slicing()

       UpperCamelCase__							:				List[Any]          =      self.get_inputs()
       pipe(**__magic_name__,		callback=__magic_name__,		callback_steps=1							)
       assert callback_fn.has_been_called
       assert number_of_steps == 3




   def       UpperCamelCase__  (    self							) ->      Optional[int]:






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

       UpperCamelCase__							:				Union[str, Any]          =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix''',		safety_checker=__magic_name__,		torch_dtype=torch.floataa							)
       UpperCamelCase__							:				List[Any]          =      pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)
       pipe.enable_attention_slicing(1							)
       pipe.enable_sequential_cpu_offload()

       UpperCamelCase__							:				List[str]          =      self.get_inputs()
       UpperCamelCase__							:				Optional[Any]          =      pipe(**__magic_name__							)

       UpperCamelCase__							:				List[str]          =      torch.cuda.max_memory_allocated()
       # make sure that less than 2.2 GB is allocated
       assert mem_bytes < 2.2 * 10**9




   def       UpperCamelCase__  (    self							) ->      int:






       """simple docstring"""
       UpperCamelCase__							:				List[Any]          =      self.get_inputs()
       # resize to resolution that is divisible by 8 but not 16 or 32
       UpperCamelCase__							:				Optional[Any]          =      inputs['''image'''].resize((504, 504)							)

       UpperCamelCase__							:				int          =      '''timbrooks/instruct-pix2pix'''
       UpperCamelCase__							:				Dict          =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           __magic_name__,		safety_checker=__magic_name__,		)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)
       pipe.enable_attention_slicing()

       UpperCamelCase__							:				Tuple          =      pipe(**__magic_name__							)
       UpperCamelCase__							:				Optional[Any]          =      output.images[0]

       UpperCamelCase__							:				Tuple          =      image[255:258, 383:386, -1]

       assert image.shape == (504, 504, 3)
       UpperCamelCase__							:				List[Any]          =      np.array([0.2726, 0.2529, 0.2664, 0.2655, 0.2641, 0.2642, 0.2591, 0.2649, 0.2590]							)

       assert np.abs(image_slice.flatten() - expected_slice							).max() < 5E-3








def 		lowerCAmelCase_    (  __UpperCAmelCase:			int      )      ->	int:
    if not isinstance(__UpperCAmelCase		,   __UpperCAmelCase      ) or number < 0:
        raise ValueError('''Input must be a non-negative integer'''      )

    UpperCamelCase__							:				Optional[Any]          =      0
    while number:
        # This way we arrive at next set bit (next 1) instead of looping
        # through each bit and checking for 1s hence the
        # loop won't run 32 times it will only run the number of `1` times
        number &= number - 1
        count += 1
    return count


if __name__ == "__main__":
  import doctest

  doctest.testmod()

