






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(			UpperCAmelCase_				,      UpperCAmelCase_				,      UpperCAmelCase_				,      unittest.TestCase   ):





       """simple docstring"""







       a							:						Union[str, Any]										=StableDiffusionInstructPixaPixPipeline
       a							:						List[str]										=TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''cross_attention_kwargs'''}
       a							:						Any										=TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
       a							:						List[str]										=IMAGE_TO_IMAGE_IMAGE_PARAMS
       a							:						Union[str, Any]										=IMAGE_TO_IMAGE_IMAGE_PARAMS





       def      _a       (   self				):
         torch.manual_seed(0				)
         UpperCamelCase_:   Tuple        =							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    ,     )
         UpperCamelCase_:   List[Any]        =							PNDMScheduler(skip_prk_steps=_lowerCamelCase				)
         torch.manual_seed(0				)
         UpperCamelCase_:   Optional[int]        =							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				)
         UpperCamelCase_:   Any        =							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    ,     )
         UpperCamelCase_:   int        =							CLIPTextModel(_lowerCamelCase				)
         UpperCamelCase_:   Union[str, Any]        =							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

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

         UpperCamelCase_:   int        =							self.get_dummy_inputs(_lowerCamelCase				)
         UpperCamelCase_:   Union[str, Any]        =							sd_pipe(**_lowerCamelCase				).images
         UpperCamelCase_:   Dict        =							image[0, -3:, -3:, -1]
         assert image.shape == (1, 3_2, 3_2, 3)
         UpperCamelCase_:   Optional[int]        =							np.array([0.7_5_2_6, 0.3_7_5_0, 0.4_5_4_7, 0.6_1_1_7, 0.5_8_6_6, 0.5_0_1_6, 0.4_3_2_7, 0.5_6_4_2, 0.4_8_1_5]				)

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





       def      _a       (   self				):
         UpperCamelCase_:   str        =							'cpu'  # ensure determinism for the device-dependent torch.Generator
         UpperCamelCase_:   int        =							self.get_dummy_components()
         UpperCamelCase_:   Dict        =							StableDiffusionInstructPixaPixPipeline(**_lowerCamelCase				)
         UpperCamelCase_:   List[Any]        =							sd_pipe.to(_lowerCamelCase				)
         sd_pipe.set_progress_bar_config(disable=_lowerCamelCase				)

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

         assert image.shape == (1, 3_2, 3_2, 3)
         UpperCamelCase_:   Dict        =							np.array([0.7_5_1_1, 0.3_6_4_2, 0.4_5_5_3, 0.6_2_3_6, 0.5_7_9_7, 0.5_0_1_3, 0.4_3_4_3, 0.5_6_1_1, 0.4_8_3_1]				)

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





       def      _a       (   self				):
         UpperCamelCase_:   int        =							'cpu'  # ensure determinism for the device-dependent torch.Generator
         UpperCamelCase_:   Tuple        =							self.get_dummy_components()
         UpperCamelCase_:   Any        =							StableDiffusionInstructPixaPixPipeline(**_lowerCamelCase				)
         UpperCamelCase_:   List[Any]        =							sd_pipe.to(_lowerCamelCase				)
         sd_pipe.set_progress_bar_config(disable=_lowerCamelCase				)

         UpperCamelCase_:   Dict        =							self.get_dummy_inputs(_lowerCamelCase				)
         UpperCamelCase_:   List[Any]        =							[inputs['prompt']] * 2

         UpperCamelCase_:   Optional[Any]        =							np.array(inputs['image']				).astype(np.floataa				) / 2_5_5.0
         UpperCamelCase_:   Any        =							torch.from_numpy(_lowerCamelCase				).unsqueeze(0				).to(_lowerCamelCase				)
         UpperCamelCase_:   Tuple        =							image / 2 + 0.5
         UpperCamelCase_:   int        =							image.permute(0    ,     3    ,     1    ,     2				)
         UpperCamelCase_:   int        =							image.repeat(2    ,     1    ,     1    ,     1				)

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

         assert image.shape == (2, 3_2, 3_2, 3)
         UpperCamelCase_:   Optional[int]        =							np.array([0.5_8_1_2, 0.5_7_4_8, 0.5_2_2_2, 0.5_9_0_8, 0.5_6_9_5, 0.7_1_7_4, 0.6_8_0_4, 0.5_5_2_3, 0.5_5_7_9]				)

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





       def      _a       (   self				):
         UpperCamelCase_:   int        =							'cpu'  # ensure determinism for the device-dependent torch.Generator
         UpperCamelCase_:   Any        =							self.get_dummy_components()
         UpperCamelCase_:   str        =							EulerAncestralDiscreteScheduler(
             beta_start=0.0_0_0_8_5    ,     beta_end=0.0_1_2    ,     beta_schedule='scaled_linear'				)
         UpperCamelCase_:   Dict        =							StableDiffusionInstructPixaPixPipeline(**_lowerCamelCase				)
         UpperCamelCase_:   str        =							sd_pipe.to(_lowerCamelCase				)
         sd_pipe.set_progress_bar_config(disable=_lowerCamelCase				)

         UpperCamelCase_:   Union[str, Any]        =							self.get_dummy_inputs(_lowerCamelCase				)
         UpperCamelCase_:   Optional[int]        =							sd_pipe(**_lowerCamelCase				).images
         UpperCamelCase_:   Optional[Any]        =							image[0, -3:, -3:, -1]

         UpperCamelCase_:   Tuple        =							[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)
         UpperCamelCase_:   Tuple        =							np.array([0.7_4_1_7, 0.3_8_4_2, 0.4_7_3_2, 0.5_7_7_6, 0.5_8_9_1, 0.5_1_3_9, 0.4_0_5_2, 0.5_6_7_3, 0.4_9_8_6]				)

         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				):
         UpperCamelCase_:   int        =							self.get_dummy_components()
         UpperCamelCase_:   Optional[Any]        =							StableDiffusionInstructPixaPixPipeline(**_lowerCamelCase				)
         UpperCamelCase_:   str        =							VaeImageProcessor(do_resize=_lowerCamelCase    ,     do_normalize=_lowerCamelCase				)
         UpperCamelCase_:   Optional[Any]        =							pipe.to(_lowerCamelCase				)
         pipe.set_progress_bar_config(disable=_lowerCamelCase				)

         UpperCamelCase_:   Dict        =							pipe(**self.get_dummy_inputs_by_type(_lowerCamelCase    ,     input_image_type='pt'				)				)[0]

         UpperCamelCase_:   Union[str, Any]        =							components['vae']
         UpperCamelCase_:   Union[str, 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():
             UpperCamelCase_:   str        =							vae.encode(inputs[image_param]				).latent_dist.mode()

         UpperCamelCase_:   Dict        =							pipe(**_lowerCamelCase				)[0]

         UpperCamelCase_:   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   ):





       """simple docstring"""







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





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

         UpperCamelCase_:   List[str]        =							self.get_inputs()
         UpperCamelCase_:   Optional[int]        =							pipe(**_lowerCamelCase				).images
         UpperCamelCase_:   Dict        =							image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         UpperCamelCase_:   Optional[Any]        =							np.array([0.5_9_0_2, 0.6_0_1_5, 0.6_0_2_7, 0.5_9_8_3, 0.6_0_9_2, 0.6_0_6_1, 0.5_7_6_5, 0.5_7_8_5, 0.5_5_5_5]				)

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





       def      _a       (   self				):
         UpperCamelCase_:   List[str]        =							StableDiffusionInstructPixaPixPipeline.from_pretrained(
             'timbrooks/instruct-pix2pix'    ,     safety_checker=_lowerCamelCase				)
         UpperCamelCase_:   List[str]        =							LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
         pipe.to(_lowerCamelCase				)
         pipe.set_progress_bar_config(disable=_lowerCamelCase				)
         pipe.enable_attention_slicing()

         UpperCamelCase_:   Union[str, Any]        =							self.get_inputs()
         UpperCamelCase_:   Optional[Any]        =							pipe(**_lowerCamelCase				).images
         UpperCamelCase_:   Optional[Any]        =							image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         UpperCamelCase_:   Any        =							np.array([0.6_5_7_8, 0.6_8_1_7, 0.6_9_7_2, 0.6_7_6_1, 0.6_8_5_6, 0.6_9_1_6, 0.6_4_2_8, 0.6_5_1_6, 0.6_3_0_1]				)

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





       def      _a       (   self				):
         UpperCamelCase_:   Optional[int]        =							StableDiffusionInstructPixaPixPipeline.from_pretrained(
             'timbrooks/instruct-pix2pix'    ,     safety_checker=_lowerCamelCase				)
         UpperCamelCase_:   Any        =							DDIMScheduler.from_config(pipe.scheduler.config				)
         pipe.to(_lowerCamelCase				)
         pipe.set_progress_bar_config(disable=_lowerCamelCase				)
         pipe.enable_attention_slicing()

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         UpperCamelCase_:   str        =							np.array([0.3_8_2_8, 0.3_8_3_4, 0.3_8_1_8, 0.3_7_9_2, 0.3_8_6_5, 0.3_7_5_2, 0.3_7_9_2, 0.3_8_4_7, 0.3_7_5_3]				)

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





       def      _a       (   self				):
         UpperCamelCase_:   Optional[Any]        =							0

         def callback_fn(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				) -> None:
           UpperCamelCase_:   List[Any]        =							True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 1:
             UpperCamelCase_:   List[str]        =							latents.detach().cpu().numpy()
             assert latents.shape == (1, 4, 6_4, 6_4)
             UpperCamelCase_:   Optional[Any]        =							latents[0, -3:, -3:, -1]
             UpperCamelCase_:   Tuple        =							np.array([-0.2_4_6_3, -0.4_6_4_4, -0.9_7_5_6, 1.5_1_7_6, 1.4_4_1_4, 0.7_8_6_6, 0.9_8_9_7, 0.8_5_2_1, 0.7_9_8_3]				)

             assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5e-2
           elif step == 2:
             UpperCamelCase_:   str        =							latents.detach().cpu().numpy()
             assert latents.shape == (1, 4, 6_4, 6_4)
             UpperCamelCase_:   Optional[int]        =							latents[0, -3:, -3:, -1]
             UpperCamelCase_:   Any        =							np.array([-0.2_6_4_4, -0.4_6_2_6, -0.9_6_5_3, 1.5_1_7_6, 1.4_5_5_1, 0.7_6_8_6, 0.9_8_0_5, 0.8_4_5_2, 0.8_1_1_5]				)

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

         UpperCamelCase_:   Tuple        =							False

         UpperCamelCase_:   Optional[int]        =							StableDiffusionInstructPixaPixPipeline.from_pretrained(
             'timbrooks/instruct-pix2pix'    ,     safety_checker=_lowerCamelCase    ,     torch_dtype=torch.floataa				)
         UpperCamelCase_:   Tuple        =							pipe.to(_lowerCamelCase				)
         pipe.set_progress_bar_config(disable=_lowerCamelCase				)
         pipe.enable_attention_slicing()

         UpperCamelCase_:   Optional[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()

         UpperCamelCase_:   List[Any]        =							StableDiffusionInstructPixaPixPipeline.from_pretrained(
             'timbrooks/instruct-pix2pix'    ,     safety_checker=_lowerCamelCase    ,     torch_dtype=torch.floataa				)
         UpperCamelCase_:   Optional[int]        =							pipe.to(_lowerCamelCase				)
         pipe.set_progress_bar_config(disable=_lowerCamelCase				)
         pipe.enable_attention_slicing(1				)
         pipe.enable_sequential_cpu_offload()

         UpperCamelCase_:   Union[str, Any]        =							self.get_inputs()
         UpperCamelCase_:   Optional[Any]        =							pipe(**_lowerCamelCase				)

         UpperCamelCase_:   str        =							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				):
         UpperCamelCase_:   Tuple        =							self.get_inputs()
         # resize to resolution that is divisible by 8 but not 16 or 32
         UpperCamelCase_:   Dict        =							inputs['image'].resize((5_0_4, 5_0_4)				)

         UpperCamelCase_:   int        =							'timbrooks/instruct-pix2pix'
         UpperCamelCase_:   Optional[Any]        =							StableDiffusionInstructPixaPixPipeline.from_pretrained(
             _lowerCamelCase    ,     safety_checker=_lowerCamelCase    ,     )
         pipe.to(_lowerCamelCase				)
         pipe.set_progress_bar_config(disable=_lowerCamelCase				)
         pipe.enable_attention_slicing()

         UpperCamelCase_:   Tuple        =							pipe(**_lowerCamelCase				)
         UpperCamelCase_:   Optional[int]        =							output.images[0]

         UpperCamelCase_:   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)
         UpperCamelCase_:   Union[str, Any]        =							np.array([0.2_7_2_6, 0.2_5_2_9, 0.2_6_6_4, 0.2_6_5_5, 0.2_6_4_1, 0.2_6_4_2, 0.2_5_9_1, 0.2_6_4_9, 0.2_5_9_0]				)

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





'''simple docstring'''







# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("""ignore""", category=UserWarning, module="""torch.optim.lr_scheduler""")







class 					_a :





     '''simple docstring'''



     def __init__(				self							,__a							,__a							,__a = True							,__a = False							)      ->				Dict:
          snake_case			:      List[Any]        =			scheduler
          snake_case			:      List[Any]        =			optimizers if isinstance(__a							,(list, tuple)							) else [optimizers]
          snake_case			:      List[str]        =			split_batches
          snake_case			:      List[str]        =			step_with_optimizer
          snake_case			:      int        =			GradientState()







     def 			snake_case_      (				self							,*__a							,**__a							)      ->				Optional[Any]:
          if not self.step_with_optimizer:
               # No link between scheduler and optimizer -> just step
               self.scheduler.step(*__a							,**__a							)
               return

          # Otherwise, first make sure the optimizer was stepped.
          if not self.gradient_state.sync_gradients:
               if self.gradient_state.adjust_scheduler:
                    self.scheduler._step_count += 1
               return

          for opt in self.optimizers:
               if opt.step_was_skipped:
                    return
          if self.split_batches:
               # Split batches -> the training dataloader batch size is not changed so one step per training step
               self.scheduler.step(*__a							,**__a							)
          else:
               # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
               # num_processes steps per training step
               snake_case			:      List[str]        =			AcceleratorState().num_processes
               for _ in range(__a							):
                    # Special case when using OneCycle and `drop_last` was not used
                    if hasattr(self.scheduler							,"""total_steps"""							):
                         if self.scheduler._step_count <= self.scheduler.total_steps:
                              self.scheduler.step(*__a							,**__a							)
                    else:
                         self.scheduler.step(*__a							,**__a							)







     def 			snake_case_      (				self							)      ->				Dict:
          return self.scheduler.get_last_lr()







     def 			snake_case_      (				self							)      ->				int:
          return self.scheduler.state_dict()







     def 			snake_case_      (				self							,__a							)      ->				Any:
          self.scheduler.load_state_dict(__a							)







     def 			snake_case_      (				self							)      ->				Any:
          return self.scheduler.get_lr()







     def 			snake_case_      (				self							,*__a							,**__a							)      ->				Optional[Any]:
          return self.scheduler.print_lr(*__a							,**__a							)



