


import argparse
import json
import os
import pickle
import shutil

import numpy as np
import torch
from distiller import Distiller
from lm_seqs_dataset import LmSeqsDataset

from transformers import (
    BertConfig,
    BertForMaskedLM,
    BertTokenizer,
    DistilBertConfig,
    DistilBertForMaskedLM,
    DistilBertTokenizer,
    GPTaConfig,
    GPTaLMHeadModel,
    GPTaTokenizer,
    RobertaConfig,
    RobertaForMaskedLM,
    RobertaTokenizer,
)
from utils import git_log, init_gpu_params, logger, set_seed


__UpperCamelCase    :						Union[str, Any]     =							{
    'distilbert': (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer),
    'roberta': (RobertaConfig, RobertaForMaskedLM, RobertaTokenizer),
    'bert': (BertConfig, BertForMaskedLM, BertTokenizer),
    'gpt2': (GPTaConfig, GPTaLMHeadModel, GPTaTokenizer),
}





def   _UpperCAmelCase       (     UpperCAmelCase      :			List[str]      ):
      """simple docstring"""




      assert (args.mlm and args.alpha_mlm > 0.0) or (not args.mlm and args.alpha_mlm == 0.0)
      assert (args.alpha_mlm > 0.0 and args.alpha_clm == 0.0) or (args.alpha_mlm == 0.0 and args.alpha_clm > 0.0)
      if args.mlm:
            assert os.path.isfile(args.token_counts      )
            assert (args.student_type in ["roberta", "distilbert"]) and (args.teacher_type in ["roberta", "bert"])
      else:
            assert (args.student_type in ["gpt2"]) and (args.teacher_type in ["gpt2"])

      assert args.teacher_type == args.student_type or (
          args.student_type == "distilbert" and args.teacher_type == "bert"
      )
      assert os.path.isfile(args.student_config      )
      if args.student_pretrained_weights is not None:
            assert os.path.isfile(args.student_pretrained_weights      )

      if args.freeze_token_type_embds:
            assert args.student_type in ["roberta"]

      assert args.alpha_ce >= 0.0
      assert args.alpha_mlm >= 0.0
      assert args.alpha_clm >= 0.0
      assert args.alpha_mse >= 0.0
      assert args.alpha_cos >= 0.0
      assert args.alpha_ce + args.alpha_mlm + args.alpha_clm + args.alpha_mse + args.alpha_cos > 0.0





def   _UpperCAmelCase       (     UpperCAmelCase      :			int						,	UpperCAmelCase      :			Optional[int]      ):
      """simple docstring"""




      if args.student_type == "roberta":
            __lowerCamelCase     :      Optional[int]       						=						False
      elif args.student_type == "gpt2":
            __lowerCamelCase     :      List[Any]       						=						False





def   _UpperCAmelCase       (     UpperCAmelCase      :			Tuple						,	UpperCAmelCase      :			Optional[int]      ):
      """simple docstring"""




      if args.student_type == "roberta":
            __lowerCamelCase     :      int       						=						False





def   _UpperCAmelCase       (     ):
      """simple docstring"""




      __lowerCamelCase     :      str       						=						argparse.ArgumentParser(description="""Training"""      )
      parser.add_argument("""--force"""						,	action="""store_true"""						,	help="""Overwrite dump_path if it already exists."""      )

      parser.add_argument(
          """--dump_path"""						,	type=UpperCAmelCase						,	required=UpperCAmelCase						,	help="""The output directory (log, checkpoints, parameters, etc.)"""      )
      parser.add_argument(
          """--data_file"""						,	type=UpperCAmelCase						,	required=UpperCAmelCase						,	help="""The binarized file (tokenized + tokens_to_ids) and grouped by sequence."""						,	)

      parser.add_argument(
          """--student_type"""						,	type=UpperCAmelCase						,	choices=["""distilbert""", """roberta""", """gpt2"""]						,	required=UpperCAmelCase						,	help="""The student type (DistilBERT, RoBERTa)."""						,	)
      parser.add_argument("""--student_config"""						,	type=UpperCAmelCase						,	required=UpperCAmelCase						,	help="""Path to the student configuration."""      )
      parser.add_argument(
          """--student_pretrained_weights"""						,	default=UpperCAmelCase						,	type=UpperCAmelCase						,	help="""Load student initialization checkpoint."""      )

      parser.add_argument(
          """--teacher_type"""						,	choices=["""bert""", """roberta""", """gpt2"""]						,	required=UpperCAmelCase						,	help="""Teacher type (BERT, RoBERTa)."""      )
      parser.add_argument("""--teacher_name"""						,	type=UpperCAmelCase						,	required=UpperCAmelCase						,	help="""The teacher model."""      )

      parser.add_argument("""--temperature"""						,	default=2.0						,	type=UpperCAmelCase						,	help="""Temperature for the softmax temperature."""      )
      parser.add_argument(
          """--alpha_ce"""						,	default=0.5						,	type=UpperCAmelCase						,	help="""Linear weight for the distillation loss. Must be >=0."""      )
      parser.add_argument(
          """--alpha_mlm"""						,	default=0.0						,	type=UpperCAmelCase						,	help="""Linear weight for the MLM loss. Must be >=0. Should be used in conjunction with `mlm` flag."""						,	)
      parser.add_argument("""--alpha_clm"""						,	default=0.5						,	type=UpperCAmelCase						,	help="""Linear weight for the CLM loss. Must be >=0."""      )
      parser.add_argument("""--alpha_mse"""						,	default=0.0						,	type=UpperCAmelCase						,	help="""Linear weight of the MSE loss. Must be >=0."""      )
      parser.add_argument(
          """--alpha_cos"""						,	default=0.0						,	type=UpperCAmelCase						,	help="""Linear weight of the cosine embedding loss. Must be >=0."""      )

      parser.add_argument(
          """--mlm"""						,	action="""store_true"""						,	help="""The LM step: MLM or CLM. If `mlm` is True, the MLM is used over CLM."""      )
      parser.add_argument(
          """--mlm_mask_prop"""						,	default=0.1_5						,	type=UpperCAmelCase						,	help="""Proportion of tokens for which we need to make a prediction."""						,	)
      parser.add_argument("""--word_mask"""						,	default=0.8						,	type=UpperCAmelCase						,	help="""Proportion of tokens to mask out."""      )
      parser.add_argument("""--word_keep"""						,	default=0.1						,	type=UpperCAmelCase						,	help="""Proportion of tokens to keep."""      )
      parser.add_argument("""--word_rand"""						,	default=0.1						,	type=UpperCAmelCase						,	help="""Proportion of tokens to randomly replace."""      )
      parser.add_argument(
          """--mlm_smoothing"""						,	default=0.7						,	type=UpperCAmelCase						,	help="""Smoothing parameter to emphasize more rare tokens (see XLM, similar to word2vec)."""						,	)
      parser.add_argument("""--token_counts"""						,	type=UpperCAmelCase						,	help="""The token counts in the data_file for MLM."""      )

      parser.add_argument(
          """--restrict_ce_to_mask"""						,	action="""store_true"""						,	help="""If true, compute the distillation loss only the [MLM] prediction distribution."""						,	)
      parser.add_argument(
          """--freeze_pos_embs"""						,	action="""store_true"""						,	help="""Freeze positional embeddings during distillation. For student_type in ['roberta', 'gpt2'] only."""						,	)
      parser.add_argument(
          """--freeze_token_type_embds"""						,	action="""store_true"""						,	help="""Freeze token type embeddings during distillation if existent. For student_type in ['roberta'] only."""						,	)

      parser.add_argument("""--n_epoch"""						,	type=UpperCAmelCase						,	default=3						,	help="""Number of pass on the whole dataset."""      )
      parser.add_argument("""--batch_size"""						,	type=UpperCAmelCase						,	default=5						,	help="""Batch size (for each process)."""      )
      parser.add_argument(
          """--group_by_size"""						,	action="""store_false"""						,	help="""If true, group sequences that have similar length into the same batch. Default is true."""						,	)

      parser.add_argument(
          """--gradient_accumulation_steps"""						,	type=UpperCAmelCase						,	default=50						,	help="""Gradient accumulation for larger training batches."""						,	)
      parser.add_argument("""--warmup_prop"""						,	default=0.0_5						,	type=UpperCAmelCase						,	help="""Linear warmup proportion."""      )
      parser.add_argument("""--weight_decay"""						,	default=0.0						,	type=UpperCAmelCase						,	help="""Weight decay if we apply some."""      )
      parser.add_argument("""--learning_rate"""						,	default=5e-4						,	type=UpperCAmelCase						,	help="""The initial learning rate for Adam."""      )
      parser.add_argument("""--adam_epsilon"""						,	default=1e-6						,	type=UpperCAmelCase						,	help="""Epsilon for Adam optimizer."""      )
      parser.add_argument("""--max_grad_norm"""						,	default=5.0						,	type=UpperCAmelCase						,	help="""Max gradient norm."""      )
      parser.add_argument("""--initializer_range"""						,	default=0.0_2						,	type=UpperCAmelCase						,	help="""Random initialization range."""      )

      parser.add_argument(
          """--fp16"""						,	action="""store_true"""						,	help="""Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"""						,	)
      parser.add_argument(
          """--fp16_opt_level"""						,	type=UpperCAmelCase						,	default="""O1"""						,	help=(
              """For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."""
              """See details at https://nvidia.github.io/apex/amp.html"""
          )						,	)
      parser.add_argument("""--n_gpu"""						,	type=UpperCAmelCase						,	default=1						,	help="""Number of GPUs in the node."""      )
      parser.add_argument("""--local_rank"""						,	type=UpperCAmelCase						,	default=-1						,	help="""Distributed training - Local rank"""      )
      parser.add_argument("""--seed"""						,	type=UpperCAmelCase						,	default=56						,	help="""Random seed"""      )

      parser.add_argument("""--log_interval"""						,	type=UpperCAmelCase						,	default=500						,	help="""Tensorboard logging interval."""      )
      parser.add_argument("""--checkpoint_interval"""						,	type=UpperCAmelCase						,	default=4_000						,	help="""Checkpoint interval."""      )
      __lowerCamelCase     :      int       						=						parser.parse_args()
      sanity_checks(UpperCAmelCase      )

      # ARGS #
      init_gpu_params(UpperCAmelCase      )
      set_seed(UpperCAmelCase      )
      if args.is_master:
            if os.path.exists(args.dump_path      ):
                  if not args.force:
                        raise ValueError(
                            f"""Serialization dir {args.dump_path} already exists, but you have not precised wheter to overwrite"""
                            """ itUse `--force` if you want to overwrite it"""      )
                  else:
                        shutil.rmtree(args.dump_path      )

            if not os.path.exists(args.dump_path      ):
                  os.makedirs(args.dump_path      )
            logger.info(f"""Experiment will be dumped and logged in {args.dump_path}"""      )

            # SAVE PARAMS #
            logger.info(f"""Param: {args}"""      )
            with open(os.path.join(args.dump_path						,	"""parameters.json"""      )						,	"""w"""      ) as f:
                  json.dump(vars(UpperCAmelCase      )						,	UpperCAmelCase						,	indent=4      )
            git_log(args.dump_path      )

      __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      Optional[int]       						=						MODEL_CLASSES[args.student_type]
      __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      List[str]       						=						MODEL_CLASSES[args.teacher_type]

      # TOKENIZER #
      __lowerCamelCase     :      Any       						=						teacher_tokenizer_class.from_pretrained(args.teacher_name      )
      __lowerCamelCase     :      Union[str, Any]       						=						{}
      for tok_name, tok_symbol in tokenizer.special_tokens_map.items():
            __lowerCamelCase     :      int       						=						tokenizer.all_special_tokens.index(UpperCAmelCase      )
            __lowerCamelCase     :      str       						=						tokenizer.all_special_ids[idx]
      logger.info(f"""Special tokens {special_tok_ids}"""      )
      __lowerCamelCase     :      Optional[Any]       						=						special_tok_ids
      __lowerCamelCase     :      int       						=						tokenizer.max_model_input_sizes[args.teacher_name]

      # DATA LOADER #
      logger.info(f"""Loading data from {args.data_file}"""      )
      with open(args.data_file						,	"""rb"""      ) as fp:
            __lowerCamelCase     :      Tuple       						=						pickle.load(UpperCAmelCase      )

      if args.mlm:
            logger.info(f"""Loading token counts from {args.token_counts} (already pre-computed)"""      )
            with open(args.token_counts						,	"""rb"""      ) as fp:
                  __lowerCamelCase     :      List[str]       						=						pickle.load(UpperCAmelCase      )

            __lowerCamelCase     :      List[str]       						=						np.maximum(UpperCAmelCase						,	1      ) ** -args.mlm_smoothing
            for idx in special_tok_ids.values():
                  __lowerCamelCase     :      List[str]       						=						0.0  # do not predict special tokens
            __lowerCamelCase     :      int       						=						torch.from_numpy(UpperCAmelCase      )
      else:
            __lowerCamelCase     :      str       						=						None

      __lowerCamelCase     :      Optional[int]       						=						LmSeqsDataset(params=UpperCAmelCase						,	data=UpperCAmelCase      )
      logger.info("""Data loader created."""      )

      # STUDENT #
      logger.info(f"""Loading student config from {args.student_config}"""      )
      __lowerCamelCase     :      List[str]       						=						student_config_class.from_pretrained(args.student_config      )
      __lowerCamelCase     :      List[Any]       						=						True

      if args.student_pretrained_weights is not None:
            logger.info(f"""Loading pretrained weights from {args.student_pretrained_weights}"""      )
            __lowerCamelCase     :      Optional[Any]       						=						student_model_class.from_pretrained(args.student_pretrained_weights						,	config=UpperCAmelCase      )
      else:
            __lowerCamelCase     :      str       						=						student_model_class(UpperCAmelCase      )

      if args.n_gpu > 0:
            student.to(f"""cuda:{args.local_rank}"""      )
      logger.info("""Student loaded."""      )

      # TEACHER #
      __lowerCamelCase     :      Any       						=						teacher_model_class.from_pretrained(args.teacher_name						,	output_hidden_states=UpperCAmelCase      )
      if args.n_gpu > 0:
            teacher.to(f"""cuda:{args.local_rank}"""      )
      logger.info(f"""Teacher loaded from {args.teacher_name}."""      )

      # FREEZING #
      if args.freeze_pos_embs:
            freeze_pos_embeddings(UpperCAmelCase						,	UpperCAmelCase      )
      if args.freeze_token_type_embds:
            freeze_token_type_embeddings(UpperCAmelCase						,	UpperCAmelCase      )

      # SANITY CHECKS #
      assert student.config.vocab_size == teacher.config.vocab_size
      assert student.config.hidden_size == teacher.config.hidden_size
      assert student.config.max_position_embeddings == teacher.config.max_position_embeddings
      if args.mlm:
            assert token_probs.size(0      ) == stu_architecture_config.vocab_size

      # DISTILLER #
      torch.cuda.empty_cache()
      __lowerCamelCase     :      Optional[Any]       						=						Distiller(
          params=UpperCAmelCase						,	dataset=UpperCAmelCase						,	token_probs=UpperCAmelCase						,	student=UpperCAmelCase						,	teacher=UpperCAmelCase      )
      distiller.train()
      logger.info("""Let's go get some drinks."""      )


if __name__ == "__main__":
    main()







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 							_UpperCamelCase (						A,unittest.TestCase						):

   '''simple docstring'''







   a_				:							Any    = "ssube/stable-diffusion-x4-upscaler-onnx"

   def   _snake_case    (  self  : Any       ,  _lowerCamelCase  : List[str]=0							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						floats_tensor((1, 3, 1_2_8, 1_2_8)       ,  rng=random.Random(_lowerCamelCase							)							)
         __lowerCamelCase     :      int       						=						torch.manual_seed(_lowerCamelCase							)
         __lowerCamelCase     :      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   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						self.get_dummy_inputs()
         __lowerCamelCase     :      Dict       						=						pipe(**_lowerCamelCase							).images
         __lowerCamelCase     :      Dict       						=						image[0, -3:, -3:, -1].flatten()

         # started as 128, should now be 512
         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      str       						=						np.array(
             [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]							)
         assert np.abs(image_slice - expected_slice							).max() < 1E-1

   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      str       						=						PNDMScheduler.from_config(pipe.scheduler.config       ,  skip_prk_steps=_lowerCamelCase							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      int       						=						np.array(
             [0.6_898_892, 0.59_240_556, 0.52_499_527, 0.58_866_215, 0.52_258_235, 0.52_572_715, 0.62_414_473, 0.6_174_387, 0.6_214_964]							)
         assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-1

   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      int       						=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      str       						=						np.array(
             [0.7_659_278, 0.76_437_664, 0.75_579_107, 0.7_691_116, 0.77_666_986, 0.7_727_672, 0.7_758_664, 0.7_812_226, 0.76_942_515]							)

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

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      Optional[Any]       						=						EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      List[str]       						=						np.array(
             [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]							)
         assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-1


   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      Union[str, Any]       						=						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      Optional[int]       						=						self.get_dummy_inputs()
         __lowerCamelCase     :      int       						=						pipe(**_lowerCamelCase							).images
         __lowerCamelCase     :      int       						=						image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      Optional[int]       						=						np.array(
             [0.77_424_496, 0.773_601, 0.7_645_288, 0.7_769_598, 0.7_772_739, 0.7_738_688, 0.78_187_233, 0.77_879_584, 0.767_043]							)

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



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

   '''simple docstring'''







   @property
   def   _snake_case    (  self  : str							):






         '''simple docstring'''



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

   @property
   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						ort.SessionOptions()
         __lowerCamelCase     :      str       						=						False
         return options

   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						load_image(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
             """/img2img/sketch-mountains-input.jpg"""							)
         __lowerCamelCase     :      Optional[Any]       						=						init_image.resize((1_2_8, 1_2_8)							)
         # using the PNDM scheduler by default
         __lowerCamelCase     :      List[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=_lowerCamelCase							)

         __lowerCamelCase     :      List[Any]       						=						"""A fantasy landscape, trending on artstation"""

         __lowerCamelCase     :      str       						=						torch.manual_seed(0							)
         __lowerCamelCase     :      Any       						=						pipe(
             prompt=_lowerCamelCase       ,  image=_lowerCamelCase       ,  guidance_scale=7.5       ,  num_inference_steps=1_0       ,  generator=_lowerCamelCase       ,  output_type="""np"""       ,  )
         __lowerCamelCase     :      List[str]       						=						output.images
         __lowerCamelCase     :      Any       						=						images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert images.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      int       						=						np.array([0.4_883, 0.4_947, 0.4_980, 0.4_975, 0.4_982, 0.4_980, 0.5_000, 0.5_006, 0.4_972]							)
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						load_image(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
             """/img2img/sketch-mountains-input.jpg"""							)
         __lowerCamelCase     :      Union[str, Any]       						=						init_image.resize((1_2_8, 1_2_8)							)
         __lowerCamelCase     :      str       						=						LMSDiscreteScheduler.from_pretrained(
             """ssube/stable-diffusion-x4-upscaler-onnx"""       ,  subfolder="""scheduler"""							)
         __lowerCamelCase     :      str       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(
             """ssube/stable-diffusion-x4-upscaler-onnx"""       ,  scheduler=_lowerCamelCase       ,  provider=self.gpu_provider       ,  sess_options=self.gpu_options       ,  )
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      str       						=						"""A fantasy landscape, trending on artstation"""

         __lowerCamelCase     :      Tuple       						=						torch.manual_seed(0							)
         __lowerCamelCase     :      Dict       						=						pipe(
             prompt=_lowerCamelCase       ,  image=_lowerCamelCase       ,  guidance_scale=7.5       ,  num_inference_steps=2_0       ,  generator=_lowerCamelCase       ,  output_type="""np"""       ,  )
         __lowerCamelCase     :      int       						=						output.images
         __lowerCamelCase     :      List[Any]       						=						images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert images.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      List[str]       						=						np.array(
             [0.50_173_753, 0.50_223_356, 0.502_039, 0.50_233_036, 0.5_023_725, 0.5_022_601, 0.5_018_758, 0.50_234_085, 0.50_241_566]							)
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



