
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__       :							int     =	StableDiffusionInstructPixaPixPipeline
     UpperCAmelCase__       :							List[Any]     =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
     UpperCAmelCase__       :							str     =	TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
     UpperCAmelCase__       :							int     =	IMAGE_TO_IMAGE_IMAGE_PARAMS
     UpperCAmelCase__       :							Optional[int]     =	IMAGE_TO_IMAGE_IMAGE_PARAMS





     def 	UpperCAmelCase(self   :					List[Any]     )		->		Optional[int]:
         torch.manual_seed(0     )
         snake_case				  =      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       ,					)
         snake_case				  =      PNDMScheduler(skip_prk_steps=lowercase_     )
         torch.manual_seed(0     )
         snake_case				  =      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     )
         snake_case				  =      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       ,					)
         snake_case				  =      CLIPTextModel(lowercase_     )
         snake_case				  =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

         snake_case				  =      {
             "unet": unet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components





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

         snake_case				  =      self.get_dummy_inputs(lowercase_     )
         snake_case				  =      sd_pipe(**lowercase_     ).images
         snake_case				  =      image[0, -3:, -3:, -1]
         assert image.shape == (1, 3_2, 3_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					Optional[Any]     )		->		Optional[int]:
         snake_case				  =      "cpu"  # ensure determinism for the device-dependent torch.Generator
         snake_case				  =      self.get_dummy_components()
         snake_case				  =      StableDiffusionInstructPixaPixPipeline(**lowercase_     )
         snake_case				  =      sd_pipe.to(lowercase_     )
         sd_pipe.set_progress_bar_config(disable=lowercase_     )

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

         assert image.shape == (1, 3_2, 3_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					Optional[Any]     )		->		Optional[int]:
         snake_case				  =      "cpu"  # ensure determinism for the device-dependent torch.Generator
         snake_case				  =      self.get_dummy_components()
         snake_case				  =      StableDiffusionInstructPixaPixPipeline(**lowercase_     )
         snake_case				  =      sd_pipe.to(lowercase_     )
         sd_pipe.set_progress_bar_config(disable=lowercase_     )

         snake_case				  =      self.get_dummy_inputs(lowercase_     )
         snake_case				  =      [inputs["prompt"]] * 2

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

         snake_case				  =      sd_pipe(**lowercase_     ).images
         snake_case				  =      image[-1, -3:, -3:, -1]

         assert image.shape == (2, 3_2, 3_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					str     )		->		Union[str, Any]:
         snake_case				  =      "cpu"  # ensure determinism for the device-dependent torch.Generator
         snake_case				  =      self.get_dummy_components()
         snake_case				  =      EulerAncestralDiscreteScheduler(
             beta_start=0.0_00_85       ,					beta_end=0.0_12       ,					beta_schedule="scaled_linear"     )
         snake_case				  =      StableDiffusionInstructPixaPixPipeline(**lowercase_     )
         snake_case				  =      sd_pipe.to(lowercase_     )
         sd_pipe.set_progress_bar_config(disable=lowercase_     )

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

         snake_case				  =      [round(lowercase_       ,					4     ) for x in image_slice.flatten().tolist()]
         print(",".join([str(lowercase_     ) for x in slice]     )     )

         assert image.shape == (1, 3_2, 3_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					Dict     )		->		Tuple:
         super().test_inference_batch_single_identical(expected_max_diff=3E-3     )





     def 	UpperCAmelCase(self   :					Union[str, Any]     )		->		Union[str, Any]:
         snake_case				  =      self.get_dummy_components()
         snake_case				  =      StableDiffusionInstructPixaPixPipeline(**lowercase_     )
         snake_case				  =      VaeImageProcessor(do_resize=lowercase_       ,					do_normalize=lowercase_     )
         snake_case				  =      pipe.to(lowercase_     )
         pipe.set_progress_bar_config(disable=lowercase_     )

         snake_case				  =      pipe(**self.get_dummy_inputs_by_type(lowercase_       ,					input_image_type="pt"     )     )[0]

         snake_case				  =      components["vae"]
         snake_case				  =      self.get_dummy_inputs_by_type(lowercase_       ,					input_image_type="pt"     )

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

         snake_case				  =      pipe(**lowercase_     )[0]

         snake_case				  =      np.abs(out - out_latents_inputs     ).max()
         self.assertLess(lowercase_       ,					1E-4       ,					"passing latents as image input generate different result from passing image"     )







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





     def 	UpperCAmelCase(self   :					Dict     )		->		Dict:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





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

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					str     )		->		List[str]:
         snake_case				  =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
             "timbrooks/instruct-pix2pix"       ,					safety_checker=lowercase_     )
         snake_case				  =      LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
         pipe.to(lowercase_     )
         pipe.set_progress_bar_config(disable=lowercase_     )
         pipe.enable_attention_slicing()

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					Tuple     )		->		Dict:
         snake_case				  =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
             "timbrooks/instruct-pix2pix"       ,					safety_checker=lowercase_     )
         snake_case				  =      DDIMScheduler.from_config(pipe.scheduler.config     )
         pipe.to(lowercase_     )
         pipe.set_progress_bar_config(disable=lowercase_     )
         pipe.enable_attention_slicing()

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         snake_case				  =      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 	UpperCAmelCase(self   :					Union[str, Any]     )		->		int:
         snake_case				  =      0

         def callback_fn(_A   :					Union[str, Any]       ,					_A   :					Dict       ,					_A   :					Optional[int]     ) -> None:
             snake_case				  =      True
             nonlocal number_of_steps
             number_of_steps += 1
             if step == 1:
                 snake_case				  =      latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 6_4, 6_4)
                 snake_case				  =      latents[0, -3:, -3:, -1]
                 snake_case				  =      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:
                 snake_case				  =      latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 6_4, 6_4)
                 snake_case				  =      latents[0, -3:, -3:, -1]
                 snake_case				  =      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

         snake_case				  =      False

         snake_case				  =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
             "timbrooks/instruct-pix2pix"       ,					safety_checker=lowercase_       ,					torch_dtype=torch.floataa     )
         snake_case				  =      pipe.to(lowercase_     )
         pipe.set_progress_bar_config(disable=lowercase_     )
         pipe.enable_attention_slicing()

         snake_case				  =      self.get_inputs()
         pipe(**lowercase_       ,					callback=lowercase_       ,					callback_steps=1     )
         assert callback_fn.has_been_called
         assert number_of_steps == 3





     def 	UpperCAmelCase(self   :					Any     )		->		Optional[Any]:
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         snake_case				  =      StableDiffusionInstructPixaPixPipeline.from_pretrained(
             "timbrooks/instruct-pix2pix"       ,					safety_checker=lowercase_       ,					torch_dtype=torch.floataa     )
         snake_case				  =      pipe.to(lowercase_     )
         pipe.set_progress_bar_config(disable=lowercase_     )
         pipe.enable_attention_slicing(1     )
         pipe.enable_sequential_cpu_offload()

         snake_case				  =      self.get_inputs()
         snake_case				  =      pipe(**lowercase_     )

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





     def 	UpperCAmelCase(self   :					Optional[int]     )		->		Optional[int]:
         snake_case				  =      self.get_inputs()
         # resize to resolution that is divisible by 8 but not 16 or 32
         snake_case				  =      inputs["image"].resize((5_0_4, 5_0_4)     )

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

         snake_case				  =      pipe(**lowercase_     )
         snake_case				  =      output.images[0]

         snake_case				  =      image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert image.shape == (5_0_4, 5_0_4, 3)
         snake_case				  =      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



import os
import unittest

from transformers.models.cpmant.tokenization_cpmant import VOCAB_FILES_NAMES, CpmAntTokenizer
from transformers.testing_utils import require_jieba, tooslow

from ...test_tokenization_common import TokenizerTesterMixin







@require_jieba
class   lowerCamelCase	(						A_							,				unittest.TestCase						):
     UpperCAmelCase__       :							Any     =	CpmAntTokenizer
     UpperCAmelCase__       :							Optional[Any]     =	False





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Dict:
         super().setUp()

         snake_case				  =      [
             "<d>",
             "</d>",
             "<s>",
             "</s>",
             "</_>",
             "<unk>",
             "<pad>",
             "</n>",
             "我",
             "是",
             "C",
             "P",
             "M",
             "A",
             "n",
             "t",
         ]
         snake_case				  =      os.path.join(self.tmpdirname       ,					VOCAB_FILES_NAMES["vocab_file"]     )
         with open(self.vocab_file       ,					"w"       ,					encoding="utf-8"     ) as vocab_writer:
             vocab_writer.write("".join([x + "\n" for x in vocab_tokens]     )     )





     @tooslow
     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Optional[int]:
         snake_case				  =      CpmAntTokenizer.from_pretrained("openbmb/cpm-ant-10b"     )
         snake_case				  =      "今天天气真好！"
         snake_case				  =      ["今天", "天气", "真", "好", "！"]
         snake_case				  =      tokenizer.tokenize(_A     )
         self.assertListEqual(_A       ,					_A     )
         snake_case				  =      "今天天气真好！"
         snake_case				  =      [tokenizer.bos_token] + tokens

         snake_case				  =      [6, 9_8_0_2, 1_4_9_6_2, 2_0_8_2, 8_3_1, 2_4_4]
         self.assertListEqual(tokenizer.convert_tokens_to_ids(_A     )       ,					_A     )

         snake_case				  =      tokenizer.decode(_A     )
         self.assertEqual(_A       ,					_A     )

