




import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    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   __a(					_a    ,						unittest.TestCase       ):






    """simple docstring"""





    lowerCAmelCase        =				'''ssube/stable-diffusion-x4-upscaler-onnx'''





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE=0  )			->		Tuple:
         UpperCAmelCase_      :							Any       =      floats_tensor((1, 3, 128, 128)       ,rng=random.Random(_SCREAMING_SNAKE_CASE  )  )
         UpperCAmelCase_      :							str       =      torch.manual_seed(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      {
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''image''': image,
             '''generator''': generator,
             '''num_inference_steps''': 3,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Union[str, Any]       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,provider='''CPUExecutionProvider'''  )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

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

         # started as 128, should now be 512
         assert image.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							Tuple       =      np.array(
             [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]  )
         assert np.abs(image_slice - expected_slice  ).max() < 1e-1





    def 							a__				(				self  )			->		str:
         UpperCAmelCase_      :							int       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,provider='''CPUExecutionProvider'''  )
         UpperCAmelCase_      :							Tuple       =      PNDMScheduler.from_config(pipe.scheduler.config       ,skip_prk_steps=_SCREAMING_SNAKE_CASE  )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

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

         assert image.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							str       =      np.array(
             [0.6_89_88_92, 0.59_24_05_56, 0.52_49_95_27, 0.58_86_62_15, 0.52_25_82_35, 0.52_57_27_15, 0.62_41_44_73, 0.6_17_43_87, 0.6_21_49_64]  )
         assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-1





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							str       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,provider='''CPUExecutionProvider'''  )
         UpperCAmelCase_      :							List[Any]       =      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

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

         assert image.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							Any       =      np.array(
             [0.7_65_92_78, 0.76_43_76_64, 0.75_57_91_07, 0.7_69_11_16, 0.77_66_69_86, 0.7_72_76_72, 0.7_75_86_64, 0.7_81_22_26, 0.76_94_25_15]  )

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





    def 							a__				(				self  )			->		Optional[int]:
         UpperCAmelCase_      :							Tuple       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,provider='''CPUExecutionProvider'''  )
         UpperCAmelCase_      :							Tuple       =      EulerDiscreteScheduler.from_config(pipe.scheduler.config  )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

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

         assert image.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							Any       =      np.array(
             [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]  )
         assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-1





    def 							a__				(				self  )			->		Dict:
         UpperCAmelCase_      :							int       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,provider='''CPUExecutionProvider'''  )
         UpperCAmelCase_      :							Tuple       =      EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config  )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

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

         assert image.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							Union[str, Any]       =      np.array(
             [0.77_42_44_96, 0.77_36_01, 0.7_64_52_88, 0.7_76_95_98, 0.7_77_27_39, 0.7_73_86_88, 0.78_18_72_33, 0.77_87_95_84, 0.76_70_43]  )

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


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






    """simple docstring"""





    @property
    def 							a__				(				self  )			->		List[Any]:
         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )





    @property
    def 							a__				(				self  )			->		Optional[int]:
         UpperCAmelCase_      :							str       =      ort.SessionOptions()
         UpperCAmelCase_      :							Any       =      False
         return options





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							List[Any]       =      load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''  )
         UpperCAmelCase_      :							Dict       =      init_image.resize((128, 128)  )
         # using the PNDM scheduler by default
         UpperCAmelCase_      :							Optional[Any]       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''       ,provider=self.gpu_provider       ,sess_options=self.gpu_options       ,)
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

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

         UpperCAmelCase_      :							Optional[int]       =      torch.manual_seed(0  )
         UpperCAmelCase_      :							int       =      pipe(
             prompt=_SCREAMING_SNAKE_CASE       ,image=_SCREAMING_SNAKE_CASE       ,guidance_scale=7.5       ,num_inference_steps=10       ,generator=_SCREAMING_SNAKE_CASE       ,output_type='''np'''       ,)
         UpperCAmelCase_      :							List[Any]       =      output.images
         UpperCAmelCase_      :							List[Any]       =      images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							Dict       =      np.array([0.48_83, 0.49_47, 0.49_80, 0.49_75, 0.49_82, 0.49_80, 0.50_00, 0.50_06, 0.49_72]  )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							Union[str, Any]       =      load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''  )
         UpperCAmelCase_      :							Optional[Any]       =      init_image.resize((128, 128)  )
         UpperCAmelCase_      :							List[str]       =      LMSDiscreteScheduler.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''       ,subfolder='''scheduler'''  )
         UpperCAmelCase_      :							str       =      OnnxStableDiffusionUpscalePipeline.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''       ,scheduler=_SCREAMING_SNAKE_CASE       ,provider=self.gpu_provider       ,sess_options=self.gpu_options       ,)
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Tuple       =      '''A fantasy landscape, trending on artstation'''

         UpperCAmelCase_      :							str       =      torch.manual_seed(0  )
         UpperCAmelCase_      :							Optional[Any]       =      pipe(
             prompt=_SCREAMING_SNAKE_CASE       ,image=_SCREAMING_SNAKE_CASE       ,guidance_scale=7.5       ,num_inference_steps=20       ,generator=_SCREAMING_SNAKE_CASE       ,output_type='''np'''       ,)
         UpperCAmelCase_      :							Any       =      output.images
         UpperCAmelCase_      :							int       =      images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 512, 3)
         UpperCAmelCase_      :							List[Any]       =      np.array(
             [0.50_17_37_53, 0.50_22_33_56, 0.50_20_39, 0.50_23_30_36, 0.5_02_37_25, 0.5_02_26_01, 0.5_01_87_58, 0.50_23_40_85, 0.50_24_15_66]  )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
__a          =			logging.getLogger(__name__)

if __name__ == "__main__":
     __a          =			argparse.ArgumentParser(
         description='Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'
     )
     parser.add_argument(
         '--data_file', type=str, default='data/dump.bert-base-uncased.pickle', help='The binarized dataset.'
     )
     parser.add_argument(
         '--token_counts_dump', type=str, default='data/token_counts.bert-base-uncased.pickle', help='The dump file.'
     )
     parser.add_argument('--vocab_size', default=30_522, type=int)
     __a          =			parser.parse_args()

     logger.info(F"""Loading data from {args.data_file}""")
     with open(args.data_file, 'rb') as fp:
          __a          =			pickle.load(fp)

     logger.info('Counting occurrences for MLM.')
     __a          =			Counter()
     for tk_ids in data:
          counter.update(tk_ids)
     __a          =			[0] * args.vocab_size
     for k, v in counter.items():
          __a          =			v

     logger.info(F"""Dump to {args.token_counts_dump}""")
     with open(args.token_counts_dump, 'wb') as handle:
          pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)