




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

@skip_mps
class 	_lowerCamelCase  (     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     unittest.TestCase				):

      """simple docstring"""

      lowerCAmelCase__     :	List[Any]      =   StableDiffusionModelEditingPipeline
      lowerCAmelCase__     :	List[Any]      =   TEXT_TO_IMAGE_PARAMS
      lowerCAmelCase__     :	Tuple      =   TEXT_TO_IMAGE_BATCH_PARAMS
      lowerCAmelCase__     :	Union[str, Any]      =   TEXT_TO_IMAGE_IMAGE_PARAMS
      lowerCAmelCase__     :	Union[str, Any]      =   TEXT_TO_IMAGE_IMAGE_PARAMS







      def  snake_case		(       self   :	Dict			):
          torch.manual_seed(0			)
          __UpperCamelCase											=  UNetaDConditionModel(
              block_out_channels=(32, 64)	,    layers_per_block=2	,    sample_size=32	,    in_channels=4	,    out_channels=4	,    down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')	,    up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')	,    cross_attention_dim=32	,    )
          __UpperCamelCase											=  DDIMScheduler()
          torch.manual_seed(0			)
          __UpperCamelCase											=  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			)
          __UpperCamelCase											=  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=1000	,    )
          __UpperCamelCase											=  CLIPTextModel(snake_case			)
          __UpperCamelCase											=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

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







      def  snake_case		(       self   :	Optional[Any]	,    snake_case   :	Union[str, Any]	,    snake_case   :	Optional[Any]=0			):
          __UpperCamelCase											=  torch.manual_seed(snake_case			)
          __UpperCamelCase											=  {
              '''prompt''': '''A field of roses''',
              '''generator''': generator,
              # Setting height and width to None to prevent OOMs on CPU.
              '''height''': None,
              '''width''': None,
              '''num_inference_steps''': 2,
              '''guidance_scale''': 6.0,
              '''output_type''': '''numpy''',
          }
          return inputs







      def  snake_case		(       self   :	Tuple			):
          __UpperCamelCase											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          __UpperCamelCase											=  self.get_dummy_components()
          __UpperCamelCase											=  StableDiffusionModelEditingPipeline(**snake_case			)
          __UpperCamelCase											=  sd_pipe.to(snake_case			)
          sd_pipe.set_progress_bar_config(disable=snake_case			)

          __UpperCamelCase											=  self.get_dummy_inputs(snake_case			)
          __UpperCamelCase											=  sd_pipe(**snake_case			).images
          __UpperCamelCase											=  image[0, -3:, -3:, -1]
          assert image.shape == (1, 64, 64, 3)

          __UpperCamelCase											=  np.array([0.47_55, 0.51_32, 0.49_76, 0.39_04, 0.35_54, 0.47_65, 0.51_39, 0.51_58, 0.48_89]			)

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







      def  snake_case		(       self   :	Union[str, Any]			):
          __UpperCamelCase											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          __UpperCamelCase											=  self.get_dummy_components()
          __UpperCamelCase											=  StableDiffusionModelEditingPipeline(**snake_case			)
          __UpperCamelCase											=  sd_pipe.to(snake_case			)
          sd_pipe.set_progress_bar_config(disable=snake_case			)

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

          assert image.shape == (1, 64, 64, 3)

          __UpperCamelCase											=  np.array([0.49_92, 0.51_01, 0.50_04, 0.39_49, 0.36_04, 0.47_35, 0.52_16, 0.52_04, 0.49_13]			)

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







      def  snake_case		(       self   :	Tuple			):
          __UpperCamelCase											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          __UpperCamelCase											=  self.get_dummy_components()
          __UpperCamelCase											=  EulerAncestralDiscreteScheduler(
              beta_start=0.0_00_85	,    beta_end=0.0_12	,    beta_schedule='''scaled_linear'''			)
          __UpperCamelCase											=  StableDiffusionModelEditingPipeline(**snake_case			)
          __UpperCamelCase											=  sd_pipe.to(snake_case			)
          sd_pipe.set_progress_bar_config(disable=snake_case			)

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

          assert image.shape == (1, 64, 64, 3)

          __UpperCamelCase											=  np.array([0.47_47, 0.53_72, 0.47_79, 0.49_82, 0.55_43, 0.48_16, 0.52_38, 0.49_04, 0.50_27]			)

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







      def  snake_case		(       self   :	int			):
          __UpperCamelCase											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          __UpperCamelCase											=  self.get_dummy_components()
          __UpperCamelCase											=  PNDMScheduler()
          __UpperCamelCase											=  StableDiffusionModelEditingPipeline(**snake_case			)
          __UpperCamelCase											=  sd_pipe.to(snake_case			)
          sd_pipe.set_progress_bar_config(disable=snake_case			)

          __UpperCamelCase											=  self.get_dummy_inputs(snake_case			)
          # the pipeline does not expect pndm so test if it raises error.
          with self.assertRaises(snake_case			):
              __UpperCamelCase											=  sd_pipe(**snake_case			).images







      def  snake_case		(       self   :	str			):
          super().test_inference_batch_single_identical(expected_max_diff=5E-3			)







      def  snake_case		(       self   :	Union[str, Any]			):
          super().test_attention_slicing_forward_pass(expected_max_diff=5E-3			)







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

      """simple docstring"""

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







      def  snake_case		(       self   :	Any	,    snake_case   :	List[str]=0			):
          __UpperCamelCase											=  torch.manual_seed(snake_case			)
          __UpperCamelCase											=  {
              '''prompt''': '''A field of roses''',
              '''generator''': generator,
              '''num_inference_steps''': 3,
              '''guidance_scale''': 7.5,
              '''output_type''': '''numpy''',
          }
          return inputs







      def  snake_case		(       self   :	Any			):
          __UpperCamelCase											=  '''CompVis/stable-diffusion-v1-4'''
          __UpperCamelCase											=  StableDiffusionModelEditingPipeline.from_pretrained(snake_case	,    safety_checker=snake_case			)
          pipe.to(snake_case			)
          pipe.set_progress_bar_config(disable=snake_case			)
          pipe.enable_attention_slicing()

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

          assert image.shape == (1, 512, 512, 3)

          __UpperCamelCase											=  np.array(
              [0.6_74_94_96, 0.6_38_64_53, 0.51_44_32_67, 0.66_09_49_05, 0.61_92_12_15, 0.5_49_13_32, 0.5_74_44_17, 0.58_07_51_06, 0.5_17_46_58]			)

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

          # make sure image changes after editing
          pipe.edit_model('''A pack of roses'''	,    '''A pack of blue roses'''			)

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

          assert image.shape == (1, 512, 512, 3)

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







      def  snake_case		(       self   :	List[str]			):
          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
          torch.cuda.reset_peak_memory_stats()

          __UpperCamelCase											=  '''CompVis/stable-diffusion-v1-4'''
          __UpperCamelCase											=  DDIMScheduler.from_pretrained(snake_case	,    subfolder='''scheduler'''			)
          __UpperCamelCase											=  StableDiffusionModelEditingPipeline.from_pretrained(
              snake_case	,    scheduler=snake_case	,    safety_checker=snake_case			)
          __UpperCamelCase											=  pipe.to(snake_case			)
          pipe.set_progress_bar_config(disable=snake_case			)
          pipe.enable_attention_slicing(1			)
          pipe.enable_sequential_cpu_offload()

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

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












import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     CpmAntConfig,
     CpmAntForCausalLM,
     CpmAntModel,
     CpmAntTokenizer,
 )

@require_torch
class 	_lowerCamelCase  :

      """simple docstring"""

      def __init__(       self   :	str	,    snake_case   :	str	,    snake_case   :	List[str]=2	,    snake_case   :	Optional[Any]=8	,    snake_case   :	List[Any]=True	,    snake_case   :	Optional[Any]=False	,    snake_case   :	str=False	,    snake_case   :	Tuple=False	,    snake_case   :	str=False	,    snake_case   :	int=99	,    snake_case   :	Tuple=32	,    snake_case   :	Union[str, Any]=3	,    snake_case   :	str=4	,    snake_case   :	List[Any]=37	,    snake_case   :	Tuple=32	,    snake_case   :	Tuple=128	,    snake_case   :	Optional[int]=8	,    snake_case   :	Optional[Any]=8	,    snake_case   :	List[Any]=8	,    snake_case   :	List[str]=1.0	,    snake_case   :	List[Any]=True	,    ):
          __UpperCamelCase											=  parent
          __UpperCamelCase											=  batch_size
          __UpperCamelCase											=  seq_length
          __UpperCamelCase											=  is_training
          __UpperCamelCase											=  use_token_type_ids
          __UpperCamelCase											=  use_input_mask
          __UpperCamelCase											=  use_labels
          __UpperCamelCase											=  use_mc_token_ids
          __UpperCamelCase											=  vocab_size
          __UpperCamelCase											=  hidden_size
          __UpperCamelCase											=  num_hidden_layers
          __UpperCamelCase											=  num_attention_heads
          __UpperCamelCase											=  intermediate_size
          __UpperCamelCase											=  num_buckets
          __UpperCamelCase											=  max_distance
          __UpperCamelCase											=  prompt_length
          __UpperCamelCase											=  prompt_types
          __UpperCamelCase											=  segment_types
          __UpperCamelCase											=  init_std
          __UpperCamelCase											=  return_dict







      def  snake_case		(       self   :	Tuple			):
          __UpperCamelCase											=  {}
          __UpperCamelCase											=  ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size			).type(torch.intaa			)
          __UpperCamelCase											=  False

          __UpperCamelCase											=  self.get_config()

          return (config, input_ids)







      def  snake_case		(       self   :	List[str]			):
          return CpmAntConfig(
              vocab_size=self.vocab_size	,    hidden_size=self.hidden_size	,    num_hidden_layers=self.num_hidden_layers	,    num_attention_heads=self.num_attention_heads	,    dim_ff=self.intermediate_size	,    position_bias_num_buckets=self.num_buckets	,    position_bias_max_distance=self.max_distance	,    prompt_types=self.prompt_types	,    prompt_length=self.prompt_length	,    segment_types=self.segment_types	,    use_cache=snake_case	,    init_std=self.init_std	,    return_dict=self.return_dict	,    )







      def  snake_case		(       self   :	Union[str, Any]	,    snake_case   :	Tuple	,    snake_case   :	List[str]	,    *snake_case   :	str			):
          __UpperCamelCase											=  CpmAntModel(config=snake_case			)
          model.to(snake_case			)
          model.eval()

          __UpperCamelCase											=  model(**snake_case			).last_hidden_state

          self.parent.assertEqual(hidden_states.shape	,    (self.batch_size, self.seq_length, config.hidden_size)			)







      def  snake_case		(       self   :	List[str]	,    snake_case   :	Optional[int]	,    snake_case   :	List[str]	,    *snake_case   :	Any			):
          __UpperCamelCase											=  CpmAntForCausalLM(snake_case			)
          model.to(snake_case			)
          __UpperCamelCase											=  input_ids['''input_ids'''].to(snake_case			)
          model.eval()

          __UpperCamelCase											=  model(**snake_case			)
          self.parent.assertEqual(
              model_output.logits.shape	,    (self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)	,    )







      def  snake_case		(       self   :	List[str]			):
          __UpperCamelCase						, __UpperCamelCase											=  self.prepare_config_and_inputs()
          return config, inputs_dict







@require_torch
class 	_lowerCamelCase  (     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     unittest.TestCase				):

      """simple docstring"""

      lowerCAmelCase__     :	List[Any]      =   (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
      lowerCAmelCase__     :	Union[str, Any]      =   (
          {"feature-extraction": CpmAntModel, "text-generation": CpmAntForCausalLM} if is_torch_available() else {}
      )

      lowerCAmelCase__     :	Optional[int]      =   False
      lowerCAmelCase__     :	Union[str, Any]      =   False
      lowerCAmelCase__     :	Optional[int]      =   False
      lowerCAmelCase__     :	Tuple      =   False
      lowerCAmelCase__     :	Union[str, Any]      =   False







      def  snake_case		(       self   :	Any			):
          __UpperCamelCase											=  CpmAntModelTester(self			)
          __UpperCamelCase											=  ConfigTester(self	,    config_class=snake_case			)







      def  snake_case		(       self   :	Optional[Any]			):
          self.config_tester.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()







      def  snake_case		(       self   :	List[str]			):
          unittest.skip('''CPMAnt doesn\'t support input_embeds.'''			)(self.test_inputs_embeds			)







      def  snake_case		(       self   :	int			):
          unittest.skip(
              '''CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'''			)(self.test_retain_grad_hidden_states_attentions			)







      def  snake_case		(       self   :	Union[str, Any]			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_cpmant_model(snake_case	,    snake_case			)







      def  snake_case		(       self   :	Any			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_lm_head_model(snake_case	,    snake_case			)







@require_torch
class 	_lowerCamelCase  (     unittest.TestCase				):

      """simple docstring"""

      @tooslow
      def  snake_case		(       self   :	Dict			):
          __UpperCamelCase											=  '''今天天气真好！'''
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntModel.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''			)
          __UpperCamelCase											=  model(**snake_case			).last_hidden_state

          __UpperCamelCase											=  torch.tensor(
              [[[6.17_08, 5.92_44, 1.08_35], [6.52_07, 6.28_93, -11.33_24], [-1.01_07, -0.05_76, -5.95_77]]]	,    )
          self.assertTrue(torch.allclose(hidden_states[:, :3, :3]	,    snake_case	,    atol=1E-2			)			)







@require_torch
class 	_lowerCamelCase  (     unittest.TestCase				):

      """simple docstring"""

      @tooslow
      def  snake_case		(       self   :	Optional[int]			):
          __UpperCamelCase											=  '''今天天气真好！'''
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntForCausalLM.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''			)
          __UpperCamelCase											=  model(**snake_case			).logits

          __UpperCamelCase											=  torch.tensor(
              [[[-6.42_67, -6.40_83, -6.39_58], [-5.88_02, -5.94_47, -5.78_11], [-5.38_96, -5.48_20, -5.42_95]]]	,    )
          self.assertTrue(torch.allclose(hidden_states[:, :3, :3]	,    snake_case	,    atol=1E-2			)			)







      @tooslow
      def  snake_case		(       self   :	List[str]			):
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntForCausalLM.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  '''今天天气不错，'''
          __UpperCamelCase											=  '''今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'''
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''			)
          __UpperCamelCase											=  model.generate(**snake_case			)
          __UpperCamelCase											=  tokenizer.batch_decode(snake_case			)
          self.assertEqual(snake_case	,    snake_case			)







      @tooslow
      def  snake_case		(       self   :	Any			):
          __UpperCamelCase											=  '''openbmb/cpm-ant-10b'''
          __UpperCamelCase											=  CpmAntForCausalLM.from_pretrained(snake_case			)
          __UpperCamelCase											=  CpmAntTokenizer.from_pretrained(snake_case			)
          __UpperCamelCase											=  ['''今天天气不错，''', '''新年快乐，万事如意！''']
          __UpperCamelCase											=  [
              '''今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的''',
              '''新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的''',
          ]
          __UpperCamelCase											=  tokenizer(snake_case	,    return_tensors='''pt'''	,    padding=snake_case			)
          __UpperCamelCase											=  model.generate(**snake_case			)
          __UpperCamelCase											=  tokenizer.batch_decode(snake_case			)
          self.assertEqual(snake_case	,    snake_case			)






