


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

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


__UpperCamelCase				:  Optional[Any]      					=							False





@skip_mps
class  lowerCAmelCase__(    snake_case__    ,    snake_case__    ,    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     str			= StableDiffusionAttendAndExcitePipeline
    A_  :     Optional[Any]			= False
    A_  :     List[str]			= TEXT_TO_IMAGE_PARAMS
    A_  :     Dict			= TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}						)
    A_  :     Optional[int]			= TEXT_TO_IMAGE_IMAGE_PARAMS
    A_  :     int			= TEXT_TO_IMAGE_IMAGE_PARAMS





    @classmethod
    def 					_lowerCamelCase (	cls				:   str	):





       '''simple docstring'''
       super().setUpClass()
       torch.use_deterministic_algorithms(__snake_case	)





    @classmethod
    def 					_lowerCamelCase (	cls				:   int	):





       '''simple docstring'''
       super().tearDownClass()
       torch.use_deterministic_algorithms(__snake_case	)





    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       torch.manual_seed(0	)
       UpperCAmelCase_							:      str						=   UNetaDConditionModel(
           block_out_channels=(32, 64)       ,   layers_per_block=1       ,   sample_size=32       ,   in_channels=4       ,   out_channels=4       ,   down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,   up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,   cross_attention_dim=32       ,   attention_head_dim=(2, 4)       ,   use_linear_projection=__snake_case       ,   )
       UpperCAmelCase_							:      Dict						=   DDIMScheduler(
           beta_start=0.00_085       ,   beta_end=0.012       ,   beta_schedule='''scaled_linear'''       ,   clip_sample=__snake_case       ,   set_alpha_to_one=__snake_case       ,   )
       torch.manual_seed(0	)
       UpperCAmelCase_							:      Optional[int]						=   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       ,   sample_size=128       ,   )
       torch.manual_seed(0	)
       UpperCAmelCase_							:      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=1_000       ,   hidden_act='''gelu'''       ,   projection_dim=512       ,   )
       UpperCAmelCase_							:      Optional[Any]						=   CLIPTextModel(__snake_case	)
       UpperCAmelCase_							:      Dict						=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)

       UpperCAmelCase_							:      List[Any]						=   {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }

       return components





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





       '''simple docstring'''
       if str(__snake_case	).startswith('''mps'''	):
          UpperCAmelCase_							:      int						=   torch.manual_seed(__snake_case	)
       else:
          UpperCAmelCase_							:      Optional[int]						=   torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
       UpperCAmelCase_							:      str						=   {
           '''prompt''': '''a cat and a frog''',
           '''token_indices''': [2, 5],
           '''generator''': generator,
           '''num_inference_steps''': 1,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
           '''max_iter_to_alter''': 2,
           '''thresholds''': {0: 0.7},
       }
       return inputs





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





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   '''cpu'''

       UpperCAmelCase_							:      List[Any]						=   self.get_dummy_components()
       UpperCAmelCase_							:      Dict						=   self.pipeline_class(**__snake_case	)
       pipe.to(__snake_case	)
       pipe.set_progress_bar_config(disable=__snake_case	)

       UpperCAmelCase_							:      Tuple						=   self.get_dummy_inputs(__snake_case	)
       UpperCAmelCase_							:      Any						=   pipe(**__snake_case	).images
       UpperCAmelCase_							:      Dict						=   image[0, -3:, -3:, -1]

       self.assertEqual(image.shape       ,   (1, 64, 64, 3)	)
       UpperCAmelCase_							:      List[Any]						=   np.array(
           [0.63_905_364, 0.62_897_307, 0.48_599_017, 0.5_133_624, 0.5_550_048, 0.45_769_516, 0.50_326_973, 0.5_023_139, 0.45_384_496]	)
       UpperCAmelCase_							:      Any						=   np.abs(image_slice.flatten() - expected_slice	).max()
       self.assertLessEqual(__snake_case       ,   1E-3	)





    def 					_lowerCamelCase (	self				:   Any	):





       '''simple docstring'''
       super().test_cpu_offload_forward_pass(expected_max_diff=5E-4	)





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
       self._test_inference_batch_consistent(batch_sizes=[1, 2]	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       self._test_inference_batch_single_identical(batch_size=2       ,   expected_max_diff=7E-4	)





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3	)





    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4	)





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       super().test_save_load_local(expected_max_difference=5E-4	)





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       super().test_save_load_optional_components(expected_max_difference=4E-4	)





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






    '''simple docstring'''





    @classmethod
    def 					_lowerCamelCase (	cls				:   Optional[Any]	):





       '''simple docstring'''
       super().setUpClass()
       torch.use_deterministic_algorithms(__snake_case	)





    @classmethod
    def 					_lowerCamelCase (	cls				:   Any	):





       '''simple docstring'''
       super().tearDownClass()
       torch.use_deterministic_algorithms(__snake_case	)





    def 					_lowerCamelCase (	self				:   Tuple	):





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





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   torch.manual_seed(51	)

       UpperCAmelCase_							:      Dict						=   StableDiffusionAttendAndExcitePipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''       ,   safety_checker=__snake_case       ,   torch_dtype=torch.floataa	)
       pipe.to('''cuda'''	)

       UpperCAmelCase_							:      Tuple						=   '''a painting of an elephant with glasses'''
       UpperCAmelCase_							:      Any						=   [5, 7]

       UpperCAmelCase_							:      Optional[Any]						=   pipe(
           prompt=__snake_case       ,   token_indices=__snake_case       ,   guidance_scale=7.5       ,   generator=__snake_case       ,   num_inference_steps=5       ,   max_iter_to_alter=5       ,   output_type='''numpy'''       ,   ).images[0]

       UpperCAmelCase_							:      Union[str, Any]						=   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy'''	)
       assert np.abs((expected_image - image).max()	) < 5E-1



import fire

from utils import calculate_rouge, save_json
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase=None			,       **__lowercase       ):
   UpperCAmelCase_							:      Tuple						=   [x.strip() for x in open(__lowercase       ).readlines()]
   UpperCAmelCase_							:      Dict						=   [x.strip() for x in open(__lowercase       ).readlines()][: len(__lowercase       )]
   UpperCAmelCase_							:      int						=   calculate_rouge(__lowercase			,       __lowercase			,       **__lowercase       )
   if save_path is not None:
      save_json(__lowercase			,       __lowercase			,       indent=__lowercase       )
   return metrics  # these print nicely


if __name__ == "__main__":
       fire.Fire(calculate_rouge_path)