


import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
       import onnxruntime as ort





class  lowerCAmelCase__(    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     str			= 'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'





    def 					_lowerCamelCase (	self				:   List[Any]       ,   __snake_case				:   Tuple=0	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   floats_tensor((1, 3, 128, 128)       ,   rng=random.Random(__snake_case	)	)
       UpperCAmelCase_							:      Dict						=   np.random.RandomState(__snake_case	)
       UpperCAmelCase_							:      List[Any]						=   {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''strength''': 0.75,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs





    def 					_lowerCamelCase (	self				:   Any	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       pipe.set_progress_bar_config(disable=__snake_case	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      List[str]						=   np.array([0.69_643, 0.58_484, 0.50_314, 0.58_760, 0.55_368, 0.59_643, 0.51_529, 0.41_217, 0.49_087]	)
       assert np.abs(image_slice - expected_slice	).max() < 1E-1





    def 					_lowerCamelCase (	self				:   Any	):





       '''simple docstring'''
       UpperCAmelCase_							:      Any						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Optional[Any]						=   PNDMScheduler.from_config(pipe.scheduler.config       ,   skip_prk_steps=__snake_case	)
       pipe.set_progress_bar_config(disable=__snake_case	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Union[str, Any]						=   np.array([0.61_737, 0.54_642, 0.53_183, 0.54_465, 0.52_742, 0.60_525, 0.49_969, 0.40_655, 0.48_154]	)

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





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      Any						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Dict						=   LMSDiscreteScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=__snake_case	)

       # warmup pass to apply optimizations
       UpperCAmelCase_							:      Dict						=   pipe(**self.get_dummy_inputs()	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      List[Any]						=   np.array([0.52_761, 0.59_977, 0.49_033, 0.49_619, 0.54_282, 0.50_311, 0.47_600, 0.40_918, 0.45_203]	)

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





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=__snake_case	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Tuple						=   np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]	)

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





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=__snake_case	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Any						=   np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]	)

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





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Dict						=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=__snake_case	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Optional[Any]						=   np.array([0.65_331, 0.58_277, 0.48_204, 0.56_059, 0.53_665, 0.56_235, 0.50_969, 0.40_009, 0.46_552]	)

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





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






    '''simple docstring'''





    @property
    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )





    @property
    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   ort.SessionOptions()
       UpperCAmelCase_							:      Tuple						=   False
       return options





    def 					_lowerCamelCase (	self				:   List[str]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/img2img/sketch-mountains-input.jpg'''	)
       UpperCAmelCase_							:      str						=   init_image.resize((768, 512)	)
       # using the PNDM scheduler by default
       UpperCAmelCase_							:      Dict						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''       ,   revision='''onnx'''       ,   safety_checker=__snake_case       ,   feature_extractor=__snake_case       ,   provider=self.gpu_provider       ,   sess_options=self.gpu_options       ,   )
       pipe.set_progress_bar_config(disable=__snake_case	)

       UpperCAmelCase_							:      List[Any]						=   '''A fantasy landscape, trending on artstation'''

       UpperCAmelCase_							:      str						=   np.random.RandomState(0	)
       UpperCAmelCase_							:      int						=   pipe(
           prompt=__snake_case       ,   image=__snake_case       ,   strength=0.75       ,   guidance_scale=7.5       ,   num_inference_steps=10       ,   generator=__snake_case       ,   output_type='''np'''       ,   )
       UpperCAmelCase_							:      Tuple						=   output.images
       UpperCAmelCase_							:      Optional[int]						=   images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 768, 3)
       UpperCAmelCase_							:      Any						=   np.array([0.4_909, 0.5_059, 0.5_372, 0.4_623, 0.4_876, 0.5_049, 0.4_820, 0.4_956, 0.5_019]	)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      str						=   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/img2img/sketch-mountains-input.jpg'''	)
       UpperCAmelCase_							:      Dict						=   init_image.resize((768, 512)	)
       UpperCAmelCase_							:      Tuple						=   LMSDiscreteScheduler.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,   subfolder='''scheduler'''       ,   revision='''onnx'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,   revision='''onnx'''       ,   scheduler=__snake_case       ,   safety_checker=__snake_case       ,   feature_extractor=__snake_case       ,   provider=self.gpu_provider       ,   sess_options=self.gpu_options       ,   )
       pipe.set_progress_bar_config(disable=__snake_case	)

       UpperCAmelCase_							:      Optional[Any]						=   '''A fantasy landscape, trending on artstation'''

       UpperCAmelCase_							:      Dict						=   np.random.RandomState(0	)
       UpperCAmelCase_							:      Any						=   pipe(
           prompt=__snake_case       ,   image=__snake_case       ,   strength=0.75       ,   guidance_scale=7.5       ,   num_inference_steps=20       ,   generator=__snake_case       ,   output_type='''np'''       ,   )
       UpperCAmelCase_							:      Optional[Any]						=   output.images
       UpperCAmelCase_							:      int						=   images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 768, 3)
       UpperCAmelCase_							:      Optional[Any]						=   np.array([0.8_043, 0.926, 0.9_581, 0.8_119, 0.8_954, 0.913, 0.7_209, 0.7_463, 0.7_431]	)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







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)