






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 				lowerCAmelCase_	(		__snake_case   ,		__snake_case   ,		__snake_case   ,		unittest.TestCase			):
				_UpperCamelCase					:       Dict							       = StableDiffusionInstructPixaPixPipeline
				_UpperCamelCase					:       Optional[int]							       = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
				_UpperCamelCase					:       Dict							       = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
				_UpperCamelCase					:       List[Any]							       = IMAGE_TO_IMAGE_IMAGE_PARAMS
				_UpperCamelCase					:       Optional[Any]							       = IMAGE_TO_IMAGE_IMAGE_PARAMS


				def 						__a					(					self      ):
											torch.manual_seed(0      )
											_lowercase							:   Optional[Any] 							=   UNetaDConditionModel(
											    block_out_channels=(3_2, 6_4) ,						layers_per_block=2 ,						sample_size=3_2 ,						in_channels=8 ,						out_channels=4 ,						down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,						up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D') ,						cross_attention_dim=3_2 ,						)
											_lowercase							:   Optional[int] 							=   PNDMScheduler(skip_prk_steps=_lowerCAmelCase      )
											torch.manual_seed(0      )
											_lowercase							:   List[Any] 							=   AutoencoderKL(
											    block_out_channels=[3_2, 6_4] ,						in_channels=3 ,						out_channels=3 ,						down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'] ,						up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'] ,						latent_channels=4 ,						)
											torch.manual_seed(0      )
											_lowercase							:   int 							=   CLIPTextConfig(
											    bos_token_id=0 ,						eos_token_id=2 ,						hidden_size=3_2 ,						intermediate_size=3_7 ,						layer_norm_eps=1E-05 ,						num_attention_heads=4 ,						num_hidden_layers=5 ,						pad_token_id=1 ,						vocab_size=1_0_0_0 ,						)
											_lowercase							:   Any 							=   CLIPTextModel(_lowerCAmelCase      )
											_lowercase							:   Tuple 							=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

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


				def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase=0      ):
											_lowercase							:   Tuple 							=   floats_tensor((1, 3, 3_2, 3_2) ,						rng=random.Random(_lowerCAmelCase      )      ).to(_lowerCAmelCase      )
											_lowercase							:   Optional[int] 							=   image.cpu().permute(0 ,						2 ,						3 ,						1      )[0]
											_lowercase							:   Dict 							=   Image.fromarray(np.uinta(_lowerCAmelCase      )      ).convert('RGB'      )
											if str(_lowerCAmelCase      ).startswith('mps'      ):
																		_lowercase							:   List[Any] 							=   torch.manual_seed(_lowerCAmelCase      )
											else:
																		_lowercase							:   int 							=   torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
											_lowercase							:   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 						__a					(					self      ):
											_lowercase							:   str 							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
											_lowercase							:   List[Any] 							=   self.get_dummy_components()
											_lowercase							:   Optional[int] 							=   StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   sd_pipe.to(_lowerCAmelCase      )
											sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   List[Any] 							=   self.get_dummy_inputs(_lowerCAmelCase      )
											_lowercase							:   Any 							=   sd_pipe(**_lowerCAmelCase      ).images
											_lowercase							:   Dict 							=   image[0, -3:, -3:, -1]
											assert image.shape == (1, 3_2, 3_2, 3)
											_lowercase							:   List[str] 							=   np.array([0.75_26, 0.37_50, 0.45_47, 0.61_17, 0.58_66, 0.50_16, 0.43_27, 0.56_42, 0.48_15]      )

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


				def 						__a					(					self      ):
											_lowercase							:   Any 							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
											_lowercase							:   int 							=   self.get_dummy_components()
											_lowercase							:   Union[str, Any] 							=   StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
											_lowercase							:   List[Any] 							=   sd_pipe.to(_lowerCAmelCase      )
											sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   Tuple 							=   self.get_dummy_inputs(_lowerCAmelCase      )
											_lowercase							:   Optional[Any] 							=   'french fries'
											_lowercase							:   Dict 							=   sd_pipe(**_lowerCAmelCase ,						negative_prompt=_lowerCAmelCase      )
											_lowercase							:   Optional[Any] 							=   output.images
											_lowercase							:   Union[str, Any] 							=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 3_2, 3_2, 3)
											_lowercase							:   Any 							=   np.array([0.75_11, 0.36_42, 0.45_53, 0.62_36, 0.57_97, 0.50_13, 0.43_43, 0.56_11, 0.48_31]      )

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


				def 						__a					(					self      ):
											_lowercase							:   Union[str, Any] 							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
											_lowercase							:   Dict 							=   self.get_dummy_components()
											_lowercase							:   int 							=   StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
											_lowercase							:   Any 							=   sd_pipe.to(_lowerCAmelCase      )
											sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   Any 							=   self.get_dummy_inputs(_lowerCAmelCase      )
											_lowercase							:   Union[str, Any] 							=   [inputs['prompt']] * 2

											_lowercase							:   Union[str, Any] 							=   np.array(inputs['image']      ).astype(np.floataa      ) / 2_55.0
											_lowercase							:   Tuple 							=   torch.from_numpy(_lowerCAmelCase      ).unsqueeze(0      ).to(_lowerCAmelCase      )
											_lowercase							:   Optional[int] 							=   image / 2 + 0.5
											_lowercase							:   List[Any] 							=   image.permute(0 ,						3 ,						1 ,						2      )
											_lowercase							:   Optional[int] 							=   image.repeat(2 ,						1 ,						1 ,						1      )

											_lowercase							:   Any 							=   sd_pipe(**_lowerCAmelCase      ).images
											_lowercase							:   List[str] 							=   image[-1, -3:, -3:, -1]

											assert image.shape == (2, 3_2, 3_2, 3)
											_lowercase							:   Optional[int] 							=   np.array([0.58_12, 0.57_48, 0.52_22, 0.59_08, 0.56_95, 0.71_74, 0.68_04, 0.55_23, 0.55_79]      )

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


				def 						__a					(					self      ):
											_lowercase							:   Optional[Any] 							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
											_lowercase							:   Dict 							=   self.get_dummy_components()
											_lowercase							:   Any 							=   EulerAncestralDiscreteScheduler(
											    beta_start=0.0_00_85 ,						beta_end=0.0_12 ,						beta_schedule='scaled_linear'      )
											_lowercase							:   Optional[Any] 							=   StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   sd_pipe.to(_lowerCAmelCase      )
											sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   Dict 							=   self.get_dummy_inputs(_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   sd_pipe(**_lowerCAmelCase      ).images
											_lowercase							:   List[Any] 							=   image[0, -3:, -3:, -1]

											_lowercase							:   Optional[int] 							=   [round(_lowerCAmelCase ,						4      ) for x in image_slice.flatten().tolist()]
											print(','.join([str(_lowerCAmelCase      ) for x in slice]      )      )

											assert image.shape == (1, 3_2, 3_2, 3)
											_lowercase							:   str 							=   np.array([0.74_17, 0.38_42, 0.47_32, 0.57_76, 0.58_91, 0.51_39, 0.40_52, 0.56_73, 0.49_86]      )

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


				def 						__a					(					self      ):
											super().test_inference_batch_single_identical(expected_max_diff=3E-3      )


				def 						__a					(					self      ):
											_lowercase							:   Union[str, Any] 							=   self.get_dummy_components()
											_lowercase							:   Dict 							=   StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   VaeImageProcessor(do_resize=_lowerCAmelCase ,						do_normalize=_lowerCAmelCase      )
											_lowercase							:   Dict 							=   pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   str 							=   pipe(**self.get_dummy_inputs_by_type(_lowerCAmelCase ,						input_image_type='pt'      )      )[0]

											_lowercase							:   List[str] 							=   components['vae']
											_lowercase							:   Optional[Any] 							=   self.get_dummy_inputs_by_type(_lowerCAmelCase ,						input_image_type='pt'      )

											for image_param in self.image_latents_params:
																		if image_param in inputs.keys():
																									_lowercase							:   Optional[int] 							=   vae.encode(inputs[image_param]      ).latent_dist.mode()

											_lowercase							:   Optional[Any] 							=   pipe(**_lowerCAmelCase      )[0]

											_lowercase							:   List[str] 							=   np.abs(out - out_latents_inputs      ).max()
											self.assertLess(_lowerCAmelCase ,						1E-4 ,						'passing latents as image input generate different result from passing image'      )




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


				def 						__a					(					self      ):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


				def 						__a					(					self ,						_lowerCAmelCase=0      ):
											_lowercase							:   Tuple 							=   torch.manual_seed(_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   load_image(
											    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'      )
											_lowercase							:   Optional[int] 							=   {
											    '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 						__a					(					self      ):
											_lowercase							:   Optional[Any] 							=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    'timbrooks/instruct-pix2pix' ,						safety_checker=_lowerCAmelCase      )
											pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
											pipe.enable_attention_slicing()

											_lowercase							:   Tuple 							=   self.get_inputs()
											_lowercase							:   Dict 							=   pipe(**_lowerCAmelCase      ).images
											_lowercase							:   Any 							=   image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											_lowercase							:   Optional[Any] 							=   np.array([0.59_02, 0.60_15, 0.60_27, 0.59_83, 0.60_92, 0.60_61, 0.57_65, 0.57_85, 0.55_55]      )

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


				def 						__a					(					self      ):
											_lowercase							:   str 							=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    'timbrooks/instruct-pix2pix' ,						safety_checker=_lowerCAmelCase      )
											_lowercase							:   Dict 							=   LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
											pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
											pipe.enable_attention_slicing()

											_lowercase							:   Optional[int] 							=   self.get_inputs()
											_lowercase							:   Optional[int] 							=   pipe(**_lowerCAmelCase      ).images
											_lowercase							:   Any 							=   image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											_lowercase							:   List[Any] 							=   np.array([0.65_78, 0.68_17, 0.69_72, 0.67_61, 0.68_56, 0.69_16, 0.64_28, 0.65_16, 0.63_01]      )

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


				def 						__a					(					self      ):
											_lowercase							:   Tuple 							=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    'timbrooks/instruct-pix2pix' ,						safety_checker=_lowerCAmelCase      )
											_lowercase							:   Optional[int] 							=   DDIMScheduler.from_config(pipe.scheduler.config      )
											pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
											pipe.enable_attention_slicing()

											_lowercase							:   Tuple 							=   self.get_inputs()
											_lowercase							:   int 							=   pipe(**_lowerCAmelCase      ).images
											_lowercase							:   Tuple 							=   image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											_lowercase							:   str 							=   np.array([0.38_28, 0.38_34, 0.38_18, 0.37_92, 0.38_65, 0.37_52, 0.37_92, 0.38_47, 0.37_53]      )

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


				def 						__a					(					self      ):
											_lowercase							:   Dict 							=   0

											def callback_fn(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ) -> None:
																		_lowercase							:   Any 							=   True
																		nonlocal number_of_steps
																		number_of_steps += 1
																		if step == 1:
																									_lowercase							:   Optional[int] 							=   latents.detach().cpu().numpy()
																									assert latents.shape == (1, 4, 6_4, 6_4)
																									_lowercase							:   Dict 							=   latents[0, -3:, -3:, -1]
																									_lowercase							:   Any 							=   np.array([-0.24_63, -0.46_44, -0.97_56, 1.51_76, 1.44_14, 0.78_66, 0.98_97, 0.85_21, 0.79_83]      )

																									assert np.abs(latents_slice.flatten() - expected_slice      ).max() < 5E-2
																		elif step == 2:
																									_lowercase							:   List[str] 							=   latents.detach().cpu().numpy()
																									assert latents.shape == (1, 4, 6_4, 6_4)
																									_lowercase							:   List[Any] 							=   latents[0, -3:, -3:, -1]
																									_lowercase							:   str 							=   np.array([-0.26_44, -0.46_26, -0.96_53, 1.51_76, 1.45_51, 0.76_86, 0.98_05, 0.84_52, 0.81_15]      )

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

											_lowercase							:   Tuple 							=   False

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

											_lowercase							:   Any 							=   self.get_inputs()
											pipe(**_lowerCAmelCase ,						callback=_lowerCAmelCase ,						callback_steps=1      )
											assert callback_fn.has_been_called
											assert number_of_steps == 3


				def 						__a					(					self      ):
											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()
											torch.cuda.reset_peak_memory_stats()

											_lowercase							:   Dict 							=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    'timbrooks/instruct-pix2pix' ,						safety_checker=_lowerCAmelCase ,						torch_dtype=torch.floataa      )
											_lowercase							:   Dict 							=   pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
											pipe.enable_attention_slicing(1      )
											pipe.enable_sequential_cpu_offload()

											_lowercase							:   List[Any] 							=   self.get_inputs()
											_lowercase							:   List[Any] 							=   pipe(**_lowerCAmelCase      )

											_lowercase							:   Optional[int] 							=   torch.cuda.max_memory_allocated()
											# make sure that less than 2.2 GB is allocated
											assert mem_bytes < 2.2 * 1_0**9


				def 						__a					(					self      ):
											_lowercase							:   Union[str, Any] 							=   self.get_inputs()
											# resize to resolution that is divisible by 8 but not 16 or 32
											_lowercase							:   Union[str, Any] 							=   inputs['image'].resize((5_0_4, 5_0_4)      )

											_lowercase							:   List[str] 							=   'timbrooks/instruct-pix2pix'
											_lowercase							:   Dict 							=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    _lowerCAmelCase ,						safety_checker=_lowerCAmelCase ,						)
											pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
											pipe.enable_attention_slicing()

											_lowercase							:   Any 							=   pipe(**_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   output.images[0]

											_lowercase							:   List[Any] 							=   image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

											assert image.shape == (5_0_4, 5_0_4, 3)
											_lowercase							:   Tuple 							=   np.array([0.27_26, 0.25_29, 0.26_64, 0.26_55, 0.26_41, 0.26_42, 0.25_91, 0.26_49, 0.25_90]      )

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










def    UpperCAmelCase_				( )   ->					list[list[int]]:
      return [list(range(10_00 - i     ,     -10_00 - i     ,     -1				)				) for i in range(10_00				)]


lowerCamelCase__					:  List[Any]      							=				generate_large_matrix()
lowerCamelCase__					:  List[Any]      							=				(
    [[4, 3, 2, -1], [3, 2, 1, -1], [1, 1, -1, -2], [-1, -1, -2, -3]],
    [[3, 2], [1, 0]],
    [[7, 7, 6]],
    [[7, 7, 6], [-1, -2, -3]],
    grid,
)

def    UpperCAmelCase_				( __UpperCAmelCase		:      list[list[int]]				)   ->					None:
      assert all(row == sorted(__UpperCAmelCase     ,     reverse=__UpperCAmelCase				) for row in grid				)
      assert all(list(__UpperCAmelCase				) == sorted(__UpperCAmelCase     ,     reverse=__UpperCAmelCase				) for col in zip(*__UpperCAmelCase				)				)

def    UpperCAmelCase_				( __UpperCAmelCase		:      list[int]				)   ->					int:
      SCREAMING_SNAKE_CASE_      		=     0
      SCREAMING_SNAKE_CASE_      		=     len(__UpperCAmelCase				) - 1

      # Edge cases such as no values or all numbers are negative.
      if not array or array[0] < 0:
            return 0

      while right + 1 > left:
            SCREAMING_SNAKE_CASE_      		=     (left + right) // 2
            SCREAMING_SNAKE_CASE_      		=     array[mid]

            # Num must be negative and the index must be greater than or equal to 0.
            if num < 0 and array[mid - 1] >= 0:
                  return mid

            if num >= 0:
                  SCREAMING_SNAKE_CASE_      		=     mid + 1
            else:
                  SCREAMING_SNAKE_CASE_      		=     mid - 1
    # No negative numbers so return the last index of the array + 1 which is the length.
      return len(__UpperCAmelCase				)

def    UpperCAmelCase_				( __UpperCAmelCase		:      list[list[int]]				)   ->					int:
      SCREAMING_SNAKE_CASE_      		=     0
      SCREAMING_SNAKE_CASE_      		=     len(grid[0]				)

      for i in range(len(__UpperCAmelCase				)				):
            SCREAMING_SNAKE_CASE_      		=     find_negative_index(grid[i][:bound]				)
            total += bound
      return (len(__UpperCAmelCase				) * len(grid[0]				)) - total

def    UpperCAmelCase_				( __UpperCAmelCase		:      list[list[int]]				)   ->					int:
      return len([number for row in grid for number in row if number < 0]				)

def    UpperCAmelCase_				( __UpperCAmelCase		:      list[list[int]]				)   ->					int:
      SCREAMING_SNAKE_CASE_      		=     0
      for row in grid:
            for i, number in enumerate(__UpperCAmelCase				):
                  if number < 0:
                        total += len(__UpperCAmelCase				) - i
                        break
      return total

def    UpperCAmelCase_				( )   ->					None:
      from timeit import timeit

      print('Running benchmarks'				)
      SCREAMING_SNAKE_CASE_      		=     (
          'from __main__ import count_negatives_binary_search, '
          'count_negatives_brute_force, count_negatives_brute_force_with_break, grid'
      )
      for func in (
          "count_negatives_binary_search",  # took 0.7727 seconds
          "count_negatives_brute_force_with_break",  # took 4.6505 seconds
          "count_negatives_brute_force",  # took 12.8160 seconds
      ):
            SCREAMING_SNAKE_CASE_      		=     timeit(f"{func}(grid=grid)"     ,     setup=__UpperCAmelCase     ,     number=5_00				)
            print(f"{func}() took {time:0.4f} seconds"				)


if __name__ == "__main__":
   import doctest

   doctest.testmod()
   benchmark()