





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



						"""simple docstring"""

						_UpperCamelCase  :	List[str]								=	StableDiffusionInstructPixaPixPipeline
						_UpperCamelCase  :	Tuple								=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
						_UpperCamelCase  :	Any								=	TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
						_UpperCamelCase  :	Union[str, Any]								=	IMAGE_TO_IMAGE_IMAGE_PARAMS
						_UpperCamelCase  :	Tuple								=	IMAGE_TO_IMAGE_IMAGE_PARAMS







						def 		__UpperCAmelCase       (    self			:		Any					):





										'''simple docstring'''







										torch.manual_seed(0					)
										_snake_case :			Optional[Any]  						=   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 ,				)
										_snake_case :			Any  						=   PNDMScheduler(skip_prk_steps=lowerCamelCase_					)
										torch.manual_seed(0					)
										_snake_case :			Optional[Any]  						=   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					)
										_snake_case :			Any  						=   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=10_00 ,				)
										_snake_case :			Tuple  						=   CLIPTextModel(lowerCamelCase_					)
										_snake_case :			List[str]  						=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

										_snake_case :			Any  						=   {
										    'unet': unet,
										    'scheduler': scheduler,
										    'vae': vae,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'safety_checker': None,
										    'feature_extractor': None,
										}
										return components







						def 		__UpperCAmelCase       (    self			:		Union[str, Any] ,				lowerCamelCase_			:		Tuple ,				lowerCamelCase_			:		List[str]=0					):





										'''simple docstring'''







										_snake_case :			List[Any]  						=   floats_tensor((1, 3, 32, 32) ,				rng=random.Random(lowerCamelCase_					)					).to(lowerCamelCase_					)
										_snake_case :			str  						=   image.cpu().permute(0 ,				2 ,				3 ,				1					)[0]
										_snake_case :			Union[str, Any]  						=   Image.fromarray(np.uinta(lowerCamelCase_					)					).convert('RGB'					)
										if str(lowerCamelCase_					).startswith('mps'					):
														_snake_case :			Any  						=   torch.manual_seed(lowerCamelCase_					)
										else:
														_snake_case :			List[Any]  						=   torch.Generator(device=lowerCamelCase_					).manual_seed(lowerCamelCase_					)
										_snake_case :			str  						=   {
										    '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			:		List[Any]					):





										'''simple docstring'''







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

										_snake_case :			int  						=   self.get_dummy_inputs(lowerCamelCase_					)
										_snake_case :			Tuple  						=   sd_pipe(**lowerCamelCase_					).images
										_snake_case :			Union[str, Any]  						=   image[0, -3:, -3:, -1]
										assert image.shape == (1, 32, 32, 3)
										_snake_case :			List[str]  						=   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			:		List[Any]					):





										'''simple docstring'''







										_snake_case :			Dict  						=   'cpu'  # ensure determinism for the device-dependent torch.Generator
										_snake_case :			List[Any]  						=   self.get_dummy_components()
										_snake_case :			Optional[int]  						=   StableDiffusionInstructPixaPixPipeline(**lowerCamelCase_					)
										_snake_case :			Optional[Any]  						=   sd_pipe.to(lowerCamelCase_					)
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase_					)

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

										assert image.shape == (1, 32, 32, 3)
										_snake_case :			List[str]  						=   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			:		Optional[Any]					):





										'''simple docstring'''







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

										_snake_case :			List[str]  						=   self.get_dummy_inputs(lowerCamelCase_					)
										_snake_case :			int  						=   [inputs['prompt']] * 2

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

										_snake_case :			Optional[int]  						=   sd_pipe(**lowerCamelCase_					).images
										_snake_case :			Any  						=   image[-1, -3:, -3:, -1]

										assert image.shape == (2, 32, 32, 3)
										_snake_case :			Any  						=   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			:		List[str]					):





										'''simple docstring'''







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

										_snake_case :			Any  						=   self.get_dummy_inputs(lowerCamelCase_					)
										_snake_case :			int  						=   sd_pipe(**lowerCamelCase_					).images
										_snake_case :			str  						=   image[0, -3:, -3:, -1]

										_snake_case :			Tuple  						=   [round(lowerCamelCase_ ,				4					) for x in image_slice.flatten().tolist()]
										print(','.join([str(lowerCamelCase_					) for x in slice]					)					)

										assert image.shape == (1, 32, 32, 3)
										_snake_case :			Any  						=   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			:		int					):





										'''simple docstring'''







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







						def 		__UpperCAmelCase       (    self			:		Optional[int]					):





										'''simple docstring'''







										_snake_case :			str  						=   self.get_dummy_components()
										_snake_case :			int  						=   StableDiffusionInstructPixaPixPipeline(**lowerCamelCase_					)
										_snake_case :			Optional[int]  						=   VaeImageProcessor(do_resize=lowerCamelCase_ ,				do_normalize=lowerCamelCase_					)
										_snake_case :			Any  						=   pipe.to(lowerCamelCase_					)
										pipe.set_progress_bar_config(disable=lowerCamelCase_					)

										_snake_case :			List[Any]  						=   pipe(**self.get_dummy_inputs_by_type(lowerCamelCase_ ,				input_image_type='pt'					)					)[0]

										_snake_case :			Optional[Any]  						=   components['vae']
										_snake_case :			Dict  						=   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():
																		_snake_case :			Optional[Any]  						=   vae.encode(inputs[image_param]					).latent_dist.mode()

										_snake_case :			Any  						=   pipe(**lowerCamelCase_					)[0]

										_snake_case :			Union[str, Any]  						=   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 			lowercase				(							unittest.TestCase     ):



						"""simple docstring"""

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





										'''simple docstring'''







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







						def 		__UpperCAmelCase       (    self			:		Optional[int] ,				lowerCamelCase_			:		int=0					):





										'''simple docstring'''







										_snake_case :			int  						=   torch.manual_seed(lowerCamelCase_					)
										_snake_case :			str  						=   load_image(
										    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'					)
										_snake_case :			Optional[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 		__UpperCAmelCase       (    self			:		int					):





										'''simple docstring'''







										_snake_case :			Any  						=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
										    'timbrooks/instruct-pix2pix' ,				safety_checker=lowerCamelCase_					)
										pipe.to(lowerCamelCase_					)
										pipe.set_progress_bar_config(disable=lowerCamelCase_					)
										pipe.enable_attention_slicing()

										_snake_case :			Tuple  						=   self.get_inputs()
										_snake_case :			Optional[Any]  						=   pipe(**lowerCamelCase_					).images
										_snake_case :			Tuple  						=   image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 5_12, 5_12, 3)
										_snake_case :			Union[str, Any]  						=   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			:		Tuple					):





										'''simple docstring'''







										_snake_case :			Optional[int]  						=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
										    'timbrooks/instruct-pix2pix' ,				safety_checker=lowerCamelCase_					)
										_snake_case :			str  						=   LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
										pipe.to(lowerCamelCase_					)
										pipe.set_progress_bar_config(disable=lowerCamelCase_					)
										pipe.enable_attention_slicing()

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

										assert image.shape == (1, 5_12, 5_12, 3)
										_snake_case :			List[Any]  						=   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'''







										_snake_case :			Tuple  						=   StableDiffusionInstructPixaPixPipeline.from_pretrained(
										    'timbrooks/instruct-pix2pix' ,				safety_checker=lowerCamelCase_					)
										_snake_case :			Tuple  						=   DDIMScheduler.from_config(pipe.scheduler.config					)
										pipe.to(lowerCamelCase_					)
										pipe.set_progress_bar_config(disable=lowerCamelCase_					)
										pipe.enable_attention_slicing()

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

										assert image.shape == (1, 5_12, 5_12, 3)
										_snake_case :			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			:		Dict					):





										'''simple docstring'''







										_snake_case :			List[str]  						=   0

										def callback_fn(lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		torch.FloatTensor					) -> None:
														_snake_case :			List[Any]  						=   True
														nonlocal number_of_steps
														number_of_steps += 1
														if step == 1:
																		_snake_case :			Dict  						=   latents.detach().cpu().numpy()
																		assert latents.shape == (1, 4, 64, 64)
																		_snake_case :			List[str]  						=   latents[0, -3:, -3:, -1]
																		_snake_case :			Union[str, Any]  						=   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:
																		_snake_case :			str  						=   latents.detach().cpu().numpy()
																		assert latents.shape == (1, 4, 64, 64)
																		_snake_case :			List[str]  						=   latents[0, -3:, -3:, -1]
																		_snake_case :			Dict  						=   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

										_snake_case :			Dict  						=   False

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

										_snake_case :			Optional[Any]  						=   self.get_inputs()
										pipe(**lowerCamelCase_ ,				callback=lowerCamelCase_ ,				callback_steps=1					)
										assert callback_fn.has_been_called
										assert number_of_steps == 3







						def 		__UpperCAmelCase       (    self			:		List[Any]					):





										'''simple docstring'''







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

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

										_snake_case :			Dict  						=   self.get_inputs()
										_snake_case :			List[str]  						=   pipe(**lowerCamelCase_					)

										_snake_case :			int  						=   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			:		Any					):





										'''simple docstring'''







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

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

										_snake_case :			Tuple  						=   pipe(**lowerCamelCase_					)
										_snake_case :			str  						=   output.images[0]

										_snake_case :			Tuple  						=   image[2_55:2_58, 3_83:3_86, -1]

										assert image.shape == (5_04, 5_04, 3)
										_snake_case :			Optional[int]  						=   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













import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class 			lowercase				(							a_		,			a_		,			a_     ):



						"""simple docstring"""

						@register_to_config
						def __init__(    self			:		Any ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		float ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		str ,				lowerCamelCase_			:		bool = False ,				):





										'''simple docstring'''







										super().__init__()

										_snake_case :			Optional[int]  						=   nn.Embedding(lowerCamelCase_ ,				lowerCamelCase_					)

										_snake_case :			List[str]  						=   nn.Embedding(lowerCamelCase_ ,				lowerCamelCase_					)
										_snake_case :			Union[str, Any]  						=   False

										_snake_case :			Tuple  						=   nn.Dropout(p=lowerCamelCase_					)

										_snake_case :			Union[str, Any]  						=   TaConfig(
										    vocab_size=lowerCamelCase_ ,				d_model=lowerCamelCase_ ,				num_heads=lowerCamelCase_ ,				d_kv=lowerCamelCase_ ,				d_ff=lowerCamelCase_ ,				dropout_rate=lowerCamelCase_ ,				feed_forward_proj=lowerCamelCase_ ,				is_decoder=lowerCamelCase_ ,				is_encoder_decoder=lowerCamelCase_ ,				)

										_snake_case :			Union[str, Any]  						=   nn.ModuleList()
										for lyr_num in range(lowerCamelCase_					):
														_snake_case :			Any  						=   TaBlock(lowerCamelCase_					)
														self.encoders.append(lowerCamelCase_					)

										_snake_case :			Tuple  						=   TaLayerNorm(lowerCamelCase_					)
										_snake_case :			List[str]  						=   nn.Dropout(p=lowerCamelCase_					)







						def 		__UpperCAmelCase       (    self			:		Optional[int] ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		Tuple					):





										'''simple docstring'''







										_snake_case :			Any  						=   self.token_embedder(lowerCamelCase_					)

										_snake_case :			List[Any]  						=   encoder_input_tokens.shape[1]
										_snake_case :			Any  						=   torch.arange(lowerCamelCase_ ,				device=encoder_input_tokens.device					)
										x += self.position_encoding(lowerCamelCase_					)

										_snake_case :			Tuple  						=   self.dropout_pre(lowerCamelCase_					)

										# inverted the attention mask
										_snake_case :			Dict  						=   encoder_input_tokens.size()
										_snake_case :			Optional[int]  						=   self.get_extended_attention_mask(lowerCamelCase_ ,				lowerCamelCase_					)

										for lyr in self.encoders:
														_snake_case :			str  						=   lyr(lowerCamelCase_ ,				lowerCamelCase_					)[0]
										_snake_case :			Any  						=   self.layer_norm(lowerCamelCase_					)

										return self.dropout_post(lowerCamelCase_					), encoder_inputs_mask






