



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       _snake_case				:






   '''simple docstring'''






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


   def       A__    (    self:     Any							)	->      str:
       UpperCAmelCase_     :       str  		=					{}
       UpperCAmelCase_     :       Union[str, Any]  		=					ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size							).type(torch.intaa							)
       UpperCAmelCase_     :       Optional[int]  		=					False

       UpperCAmelCase_     :       Dict  		=					self.get_config()

       return (config, input_ids)


   def       A__    (    self:     int							)	->      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=lowerCamelCase_  ,init_std=self.init_std  ,return_dict=self.return_dict  ,)


   def       A__    (    self:     Union[str, Any]  ,lowerCamelCase_:     Dict  ,lowerCamelCase_:     Any  ,*lowerCamelCase_:     Optional[int]							)	->      List[str]:
       UpperCAmelCase_     :       Optional[Any]  		=					CpmAntModel(config=lowerCamelCase_							)
       model.to(lowerCamelCase_							)
       model.eval()

       UpperCAmelCase_     :       int  		=					model(**lowerCamelCase_							).last_hidden_state

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


   def       A__    (    self:     Optional[Any]  ,lowerCamelCase_:     List[Any]  ,lowerCamelCase_:     Optional[int]  ,*lowerCamelCase_:     List[Any]							)	->      str:
       UpperCAmelCase_     :       str  		=					CpmAntForCausalLM(lowerCamelCase_							)
       model.to(lowerCamelCase_							)
       UpperCAmelCase_     :       List[Any]  		=					input_ids["""input_ids"""].to(lowerCamelCase_							)
       model.eval()

       UpperCAmelCase_     :       str  		=					model(**lowerCamelCase_							)
       self.parent.assertEqual(
           model_output.logits.shape  ,(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)  ,)






   def       A__    (    self:     Dict							)	->      Union[str, Any]:
       UpperCAmelCase_		, UpperCAmelCase_     :       int  		=					self.prepare_config_and_inputs()
       return config, inputs_dict







@require_torch
class       _snake_case				(     __snake_case ,							__snake_case ,							unittest.TestCase ):






   '''simple docstring'''






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

   A__      : int       =						False
   A__      : Any       =						False
   A__      : Tuple       =						False
   A__      : Tuple       =						False
   A__      : List[Any]       =						False


   def       A__    (    self:     int							)	->      Union[str, Any]:
       UpperCAmelCase_     :       Union[str, Any]  		=					CpmAntModelTester(self							)
       UpperCAmelCase_     :       List[str]  		=					ConfigTester(self  ,config_class=lowerCamelCase_							)


   def       A__    (    self:     Dict							)	->      Dict:
       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       A__    (    self:     str							)	->      int:
       unittest.skip("""CPMAnt doesn't support input_embeds."""							)(self.test_inputs_embeds							)


   def       A__    (    self:     Optional[int]							)	->      List[str]:
       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       A__    (    self:     Union[str, Any]							)	->      List[Any]:
       UpperCAmelCase_		, UpperCAmelCase_     :       List[str]  		=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_cpmant_model(lowerCamelCase_  ,lowerCamelCase_							)






   def       A__    (    self:     int							)	->      int:
       UpperCAmelCase_		, UpperCAmelCase_     :       Optional[Any]  		=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_lm_head_model(lowerCamelCase_  ,lowerCamelCase_							)







@require_torch
class       _snake_case				(     unittest.TestCase ):






   '''simple docstring'''






   @tooslow
   def       A__    (    self:     List[str]							)	->      Tuple:
       UpperCAmelCase_     :       str  		=					"""今天天气真好！"""
       UpperCAmelCase_     :       Dict  		=					"""openbmb/cpm-ant-10b"""
       UpperCAmelCase_     :       Dict  		=					CpmAntModel.from_pretrained(lowerCamelCase_							)
       UpperCAmelCase_     :       int  		=					CpmAntTokenizer.from_pretrained(lowerCamelCase_							)
       UpperCAmelCase_     :       str  		=					tokenizer(lowerCamelCase_  ,return_tensors="""pt"""							)
       UpperCAmelCase_     :       Tuple  		=					model(**lowerCamelCase_							).last_hidden_state

       UpperCAmelCase_     :       List[Any]  		=					torch.tensor(
           [[[6.1_7_0_8, 5.9_2_4_4, 1.0_8_3_5], [6.5_2_0_7, 6.2_8_9_3, -1_1.3_3_2_4], [-1.0_1_0_7, -0.0_5_7_6, -5.9_5_7_7]]]  ,)
       self.assertTrue(torch.allclose(hidden_states[:, :3, :3]  ,lowerCamelCase_  ,atol=1e-2							)							)







@require_torch
class       _snake_case				(     unittest.TestCase ):






   '''simple docstring'''






   @tooslow
   def       A__    (    self:     List[str]							)	->      Any:
       UpperCAmelCase_     :       Optional[Any]  		=					"""今天天气真好！"""
       UpperCAmelCase_     :       Optional[int]  		=					"""openbmb/cpm-ant-10b"""
       UpperCAmelCase_     :       Union[str, Any]  		=					CpmAntForCausalLM.from_pretrained(lowerCamelCase_							)
       UpperCAmelCase_     :       Any  		=					CpmAntTokenizer.from_pretrained(lowerCamelCase_							)
       UpperCAmelCase_     :       str  		=					tokenizer(lowerCamelCase_  ,return_tensors="""pt"""							)
       UpperCAmelCase_     :       Dict  		=					model(**lowerCamelCase_							).logits

       UpperCAmelCase_     :       Any  		=					torch.tensor(
           [[[-6.4_2_6_7, -6.4_0_8_3, -6.3_9_5_8], [-5.8_8_0_2, -5.9_4_4_7, -5.7_8_1_1], [-5.3_8_9_6, -5.4_8_2_0, -5.4_2_9_5]]]  ,)
       self.assertTrue(torch.allclose(hidden_states[:, :3, :3]  ,lowerCamelCase_  ,atol=1e-2							)							)


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






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




import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class 			lowerCAmelCase__(     snake_case__     ,     unittest.TestCase    ):






    '''simple docstring'''





    A_	:  Tuple								=			RePaintPipeline
    A_	:  Tuple								=			IMAGE_INPAINTING_PARAMS - {'width', 'height', 'guidance_scale'}
    A_	:  str								=			PipelineTesterMixin.required_optional_params - {
        'latents',
        'num_images_per_prompt',
        'callback',
        'callback_steps',
    }
    A_	:  Optional[Any]								=			IMAGE_INPAINTING_BATCH_PARAMS
    A_	:  Union[str, Any]								=			False





    def        _lowerCamelCase			(						self					: Any						):





       '''simple docstring'''
       torch.manual_seed(0						)
       torch.manual_seed(0						)
       UpperCAmelCase_				:   int							= UNetaDModel(
           block_out_channels=(32, 64)							, layers_per_block=2							, sample_size=32							, in_channels=3							, out_channels=3							, down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')							, up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')							, )
       UpperCAmelCase_				:   str							= RePaintScheduler()
       UpperCAmelCase_				:   Optional[Any]							= {'''unet''': unet, '''scheduler''': scheduler}
       return components





    def        _lowerCamelCase			(						self					: List[str]							, __snake_case					: Optional[Any]							, __snake_case					: int=0						):





       '''simple docstring'''
       if str(__snake_case						).startswith('''mps'''						):
          UpperCAmelCase_				:   Optional[Any]							= torch.manual_seed(__snake_case						)
       else:
          UpperCAmelCase_				:   Optional[Any]							= torch.Generator(device=__snake_case						).manual_seed(__snake_case						)
       UpperCAmelCase_				:   Optional[Any]							= np.random.RandomState(__snake_case						).standard_normal((1, 3, 32, 32)						)
       UpperCAmelCase_				:   Optional[int]							= torch.from_numpy(__snake_case						).to(device=__snake_case							, dtype=torch.floataa						)
       UpperCAmelCase_				:   Any							= (image > 0).to(device=__snake_case							, dtype=torch.floataa						)
       UpperCAmelCase_				:   List[str]							= {
           '''image''': image,
           '''mask_image''': mask,
           '''generator''': generator,
           '''num_inference_steps''': 5,
           '''eta''': 0.0,
           '''jump_length''': 2,
           '''jump_n_sample''': 2,
           '''output_type''': '''numpy''',
       }
       return inputs





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





       '''simple docstring'''
       UpperCAmelCase_				:   str							= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase_				:   Optional[int]							= self.get_dummy_components()
       UpperCAmelCase_				:   Dict							= RePaintPipeline(**__snake_case						)
       UpperCAmelCase_				:   int							= sd_pipe.to(__snake_case						)
       sd_pipe.set_progress_bar_config(disable=__snake_case						)

       UpperCAmelCase_				:   List[Any]							= self.get_dummy_inputs(__snake_case						)
       UpperCAmelCase_				:   Union[str, Any]							= sd_pipe(**__snake_case						).images
       UpperCAmelCase_				:   Optional[Any]							= image[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       UpperCAmelCase_				:   Tuple							= np.array([1.0_000, 0.5_426, 0.5_497, 0.2_200, 1.0_000, 1.0_000, 0.5_623, 1.0_000, 0.6_274]						)

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





    @skip_mps
    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       return super().test_save_load_local()





    @unittest.skip('''non-deterministic pipeline'''						)
    def        _lowerCamelCase			(						self					: Optional[Any]						):





       '''simple docstring'''
       return super().test_inference_batch_single_identical()





    @skip_mps
    def        _lowerCamelCase			(						self					: Dict						):





       '''simple docstring'''
       return super().test_dict_tuple_outputs_equivalent()





    @skip_mps
    def        _lowerCamelCase			(						self					: int						):





       '''simple docstring'''
       return super().test_save_load_optional_components()





    @skip_mps
    def        _lowerCamelCase			(						self					: Union[str, Any]						):





       '''simple docstring'''
       return super().test_attention_slicing_forward_pass()





@nightly
@require_torch_gpu
class 			lowerCAmelCase__(     unittest.TestCase    ):






    '''simple docstring'''





    def        _lowerCamelCase			(						self					: Tuple						):





       '''simple docstring'''
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





    def        _lowerCamelCase			(						self					: List[str]						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[Any]							= load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'''
           '''repaint/celeba_hq_256.png'''						)
       UpperCAmelCase_				:   int							= load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png'''						)
       UpperCAmelCase_				:   List[Any]							= load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'''
           '''repaint/celeba_hq_256_result.npy'''						)

       UpperCAmelCase_				:   List[str]							= '''google/ddpm-ema-celebahq-256'''
       UpperCAmelCase_				:   Union[str, Any]							= UNetaDModel.from_pretrained(__snake_case						)
       UpperCAmelCase_				:   Dict							= RePaintScheduler.from_pretrained(__snake_case						)

       UpperCAmelCase_				:   Optional[int]							= RePaintPipeline(unet=__snake_case							, scheduler=__snake_case						).to(__snake_case						)
       repaint.set_progress_bar_config(disable=__snake_case						)
       repaint.enable_attention_slicing()

       UpperCAmelCase_				:   Optional[int]							= torch.manual_seed(0						)
       UpperCAmelCase_				:   List[str]							= repaint(
           __snake_case							, __snake_case							, num_inference_steps=250							, eta=0.0							, jump_length=10							, jump_n_sample=10							, generator=__snake_case							, output_type='''np'''							, )
       UpperCAmelCase_				:   Tuple							= output.images[0]

       assert image.shape == (256, 256, 3)
       assert np.abs(expected_image - image						).mean() < 1E-2