




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     SCREAMING_SNAKE_CASE__		(_a    ,					_a    ,					_a    ,					unittest.TestCase			):
     lowercase_ :      Any				= StableDiffusionInstructPixaPixPipeline
     lowercase_ :      Optional[Any]				= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''cross_attention_kwargs'''}
     lowercase_ :      Any				= TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
     lowercase_ :      List[str]				= IMAGE_TO_IMAGE_IMAGE_PARAMS
     lowercase_ :      str				= IMAGE_TO_IMAGE_IMAGE_PARAMS
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       torch.manual_seed(0						)
       lowerCAmelCase__										=      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       ,						)
       lowerCAmelCase__										=      PNDMScheduler(skip_prk_steps=A__						)
       torch.manual_seed(0						)
       lowerCAmelCase__										=      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						)
       lowerCAmelCase__										=      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       ,						)
       lowerCAmelCase__										=      CLIPTextModel(A__						)
       lowerCAmelCase__										=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

       lowerCAmelCase__										=      {
           """unet""": unet,
           """scheduler""": scheduler,
           """vae""": vae,
           """text_encoder""": text_encoder,
           """tokenizer""": tokenizer,
           """safety_checker""": None,
           """feature_extractor""": None,
       }
       return components
     def        A__      (						self	:		List[str]       ,						__lowerCamelCase	:		Union[str, Any]       ,						__lowerCamelCase	:		str=0						):



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



       """simple docstring"""
       lowerCAmelCase__										=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
       lowerCAmelCase__										=      self.get_dummy_components()
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline(**A__						)
       lowerCAmelCase__										=      sd_pipe.to(A__						)
       sd_pipe.set_progress_bar_config(disable=A__						)

       lowerCAmelCase__										=      self.get_dummy_inputs(A__						)
       lowerCAmelCase__										=      sd_pipe(**A__						).images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1]
       assert image.shape == (1, 32, 32, 3)
       lowerCAmelCase__										=      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        A__      (						self	:		Any						):



       """simple docstring"""
       lowerCAmelCase__										=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
       lowerCAmelCase__										=      self.get_dummy_components()
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline(**A__						)
       lowerCAmelCase__										=      sd_pipe.to(A__						)
       sd_pipe.set_progress_bar_config(disable=A__						)

       lowerCAmelCase__										=      self.get_dummy_inputs(A__						)
       lowerCAmelCase__										=      """french fries"""
       lowerCAmelCase__										=      sd_pipe(**A__       ,						negative_prompt=A__						)
       lowerCAmelCase__										=      output.images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1]

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



       """simple docstring"""
       lowerCAmelCase__										=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
       lowerCAmelCase__										=      self.get_dummy_components()
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline(**A__						)
       lowerCAmelCase__										=      sd_pipe.to(A__						)
       sd_pipe.set_progress_bar_config(disable=A__						)

       lowerCAmelCase__										=      self.get_dummy_inputs(A__						)
       lowerCAmelCase__										=      [inputs["""prompt"""]] * 2

       lowerCAmelCase__										=      np.array(inputs['''image''']						).astype(np.floataa						) / 2_55.0
       lowerCAmelCase__										=      torch.from_numpy(A__						).unsqueeze(0						).to(A__						)
       lowerCAmelCase__										=      image / 2 + 0.5
       lowerCAmelCase__										=      image.permute(0       ,						3       ,						1       ,						2						)
       lowerCAmelCase__										=      image.repeat(2       ,						1       ,						1       ,						1						)

       lowerCAmelCase__										=      sd_pipe(**A__						).images
       lowerCAmelCase__										=      image[-1, -3:, -3:, -1]

       assert image.shape == (2, 32, 32, 3)
       lowerCAmelCase__										=      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        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
       lowerCAmelCase__										=      self.get_dummy_components()
       lowerCAmelCase__										=      EulerAncestralDiscreteScheduler(
           beta_start=0.0_0085       ,						beta_end=0.012       ,						beta_schedule='''scaled_linear'''						)
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline(**A__						)
       lowerCAmelCase__										=      sd_pipe.to(A__						)
       sd_pipe.set_progress_bar_config(disable=A__						)

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

       lowerCAmelCase__										=      [round(A__       ,						4						) for x in image_slice.flatten().tolist()]
       print(''','''.join([str(A__						) for x in slice]						)						)

       assert image.shape == (1, 32, 32, 3)
       lowerCAmelCase__										=      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        A__      (						self	:		str						):



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



     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.get_dummy_components()
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline(**A__						)
       lowerCAmelCase__										=      VaeImageProcessor(do_resize=A__       ,						do_normalize=A__						)
       lowerCAmelCase__										=      pipe.to(A__						)
       pipe.set_progress_bar_config(disable=A__						)

       lowerCAmelCase__										=      pipe(**self.get_dummy_inputs_by_type(A__       ,						input_image_type='''pt'''						)						)[0]

       lowerCAmelCase__										=      components["""vae"""]
       lowerCAmelCase__										=      self.get_dummy_inputs_by_type(A__       ,						input_image_type='''pt'''						)

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

       lowerCAmelCase__										=      pipe(**A__						)[0]

       lowerCAmelCase__										=      np.abs(out - out_latents_inputs						).max()
       self.assertLess(A__       ,						1e-4       ,						'''passing latents as image input generate different result from passing image'''						)





@slow
@require_torch_gpu
class     SCREAMING_SNAKE_CASE__		(unittest.TestCase			):
     def        A__      (						self	:		Union[str, Any]						):



       """simple docstring"""
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
     def        A__      (						self	:		Optional[int]       ,						__lowerCamelCase	:		Union[str, Any]=0						):



       """simple docstring"""
       lowerCAmelCase__										=      torch.manual_seed(A__						)
       lowerCAmelCase__										=      load_image(
           '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'''						)
       lowerCAmelCase__										=      {
           """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	:		Dict						):



       """simple docstring"""
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix'''       ,						safety_checker=A__						)
       pipe.to(A__						)
       pipe.set_progress_bar_config(disable=A__						)
       pipe.enable_attention_slicing()

       lowerCAmelCase__										=      self.get_inputs()
       lowerCAmelCase__										=      pipe(**A__						).images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 5_12, 5_12, 3)
       lowerCAmelCase__										=      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        A__      (						self	:		int						):



       """simple docstring"""
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix'''       ,						safety_checker=A__						)
       lowerCAmelCase__										=      LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.to(A__						)
       pipe.set_progress_bar_config(disable=A__						)
       pipe.enable_attention_slicing()

       lowerCAmelCase__										=      self.get_inputs()
       lowerCAmelCase__										=      pipe(**A__						).images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1].flatten()

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



       """simple docstring"""
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix'''       ,						safety_checker=A__						)
       lowerCAmelCase__										=      DDIMScheduler.from_config(pipe.scheduler.config						)
       pipe.to(A__						)
       pipe.set_progress_bar_config(disable=A__						)
       pipe.enable_attention_slicing()

       lowerCAmelCase__										=      self.get_inputs()
       lowerCAmelCase__										=      pipe(**A__						).images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 5_12, 5_12, 3)
       lowerCAmelCase__										=      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        A__      (						self	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      0

       def callback_fn(__lowerCamelCase	:		Tuple       ,						__lowerCamelCase	:		Tuple       ,						__lowerCamelCase	:		Tuple						) -> None:
         lowerCAmelCase__										=      True
         nonlocal number_of_steps
         number_of_steps += 1
         if step == 1:
           lowerCAmelCase__										=      latents.detach().cpu().numpy()
           assert latents.shape == (1, 4, 64, 64)
           lowerCAmelCase__										=      latents[0, -3:, -3:, -1]
           lowerCAmelCase__										=      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:
           lowerCAmelCase__										=      latents.detach().cpu().numpy()
           assert latents.shape == (1, 4, 64, 64)
           lowerCAmelCase__										=      latents[0, -3:, -3:, -1]
           lowerCAmelCase__										=      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

       lowerCAmelCase__										=      False

       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           '''timbrooks/instruct-pix2pix'''       ,						safety_checker=A__       ,						torch_dtype=torch.floataa						)
       lowerCAmelCase__										=      pipe.to(A__						)
       pipe.set_progress_bar_config(disable=A__						)
       pipe.enable_attention_slicing()

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



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

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

       lowerCAmelCase__										=      self.get_inputs()
       lowerCAmelCase__										=      pipe(**A__						)

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



     def        A__      (						self	:		Dict						):



       """simple docstring"""
       lowerCAmelCase__										=      self.get_inputs()
       # resize to resolution that is divisible by 8 but not 16 or 32
       lowerCAmelCase__										=      inputs["""image"""].resize((5_04, 5_04)						)

       lowerCAmelCase__										=      """timbrooks/instruct-pix2pix"""
       lowerCAmelCase__										=      StableDiffusionInstructPixaPixPipeline.from_pretrained(
           A__       ,						safety_checker=A__       ,						)
       pipe.to(A__						)
       pipe.set_progress_bar_config(disable=A__						)
       pipe.enable_attention_slicing()

       lowerCAmelCase__										=      pipe(**A__						)
       lowerCAmelCase__										=      output.images[0]

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

       assert image.shape == (5_04, 5_04, 3)
       lowerCAmelCase__										=      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 unittest

from transformers import load_tool
from transformers.utils import is_torch_available


if is_torch_available():
	import torch

from transformers.testing_utils import require_torch

from .test_tools_common import ToolTesterMixin



@require_torch
class        _a		(unittest.TestCase						,  __magic_name__							):


					'''simple docstring'''






					def     __A     (					self     ):
						A__				:       List[str]	  =						load_tool("""text-to-speech"""     )
						self.tool.setup()






					def     __A     (					self     ):
						# SpeechT5 isn't deterministic
						torch.manual_seed(0     )
						A__				:       Tuple	  =						self.tool("""hey"""     )
						A__				:       Tuple	  =						result.to_raw()
						self.assertTrue(
						    torch.allclose(
						        resulting_tensor[:3]       ,			torch.tensor([-0.0_0_0_5_9_6_6_6_6_8_8_3_2_1_1_5_8_2_9, -0.0_0_0_3_6_5_7_6_4_0_1_9_0_7_9_5_0_6_4, -0.0_0_0_1_3_4_3_9_5_0_2_7_9_9_8_8_3_4_8_5]     )       ,			)     )






					def     __A     (					self     ):
						# SpeechT5 isn't deterministic
						torch.manual_seed(0     )
						A__				:       Any	  =						self.tool("""hey"""     )
						A__				:       Optional[Any]	  =						result.to_raw()
						self.assertTrue(
						    torch.allclose(
						        resulting_tensor[:3]       ,			torch.tensor([-0.0_0_0_5_9_6_6_6_6_8_8_3_2_1_1_5_8_2_9, -0.0_0_0_3_6_5_7_6_4_0_1_9_0_7_9_5_0_6_4, -0.0_0_0_1_3_4_3_9_5_0_2_7_9_9_8_8_3_4_8_5]     )       ,			)     )



