



"""simple docstring"""






import argparse
import json
import os
import re
from collections import OrderedDict
from os.path import basename, dirname

import fairseq
import torch
from fairseq import hub_utils
from fairseq.data.dictionary import Dictionary

from transformers import FSMTConfig, FSMTForConditionalGeneration
from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES
from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
from transformers.utils import WEIGHTS_NAME, logging


logging.set_verbosity_warning()

lowerCAmelCase__      =    2

# based on the results of a search on a range of `num_beams`, `length_penalty` and `early_stopping`
# values against wmt19 test data to obtain the best BLEU scores, we will use the following defaults:
#
# * `num_beams`: 5 (higher scores better, but requires more memory/is slower, can be adjusted by users)
# * `early_stopping`: `False` consistently scored better
# * `length_penalty` varied, so will assign the best one depending on the model
lowerCAmelCase__      =    {
    # fairseq:
    '''wmt19-ru-en''': {'''length_penalty''': 1.1},
    '''wmt19-en-ru''': {'''length_penalty''': 1.1_5},
    '''wmt19-en-de''': {'''length_penalty''': 1.0},
    '''wmt19-de-en''': {'''length_penalty''': 1.1},
    # allenai:
    '''wmt16-en-de-dist-12-1''': {'''length_penalty''': 0.6},
    '''wmt16-en-de-dist-6-1''': {'''length_penalty''': 0.6},
    '''wmt16-en-de-12-1''': {'''length_penalty''': 0.8},
    '''wmt19-de-en-6-6-base''': {'''length_penalty''': 0.6},
    '''wmt19-de-en-6-6-big''': {'''length_penalty''': 0.6},
}

# this remaps the different models to their organization names
lowerCAmelCase__      =    {}
for m in ["wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"]:
   lowerCAmelCase__      =    '''facebook'''
for m in [
    "wmt16-en-de-dist-12-1",
    "wmt16-en-de-dist-6-1",
    "wmt16-en-de-12-1",
    "wmt19-de-en-6-6-base",
    "wmt19-de-en-6-6-big",
]:
   lowerCAmelCase__      =    '''allenai'''






def 	snake_case_  (				A_ :		List[Any]						):


       '''simple docstring'''





       _lowerCamelCase						:    List[Any]         =			dict((re.sub(R'''@@$''',    '''''',    A_						), v) if k.endswith('''@@'''						) else (re.sub(R'''$''',    '''</w>''',    A_						), v) for k, v in d.items()						)
       _lowerCamelCase						:    str         =			'''<s> <pad> </s> <unk>'''.split()
       # restore the special tokens
       for k in keep_keys:
              del da[F'''{k}</w>''']
              _lowerCamelCase						:    int         =			d[k]  # restore
       return da






def 	snake_case_  (				A_ :		List[Any],    A_ :		Union[str, Any]						):


       '''simple docstring'''





       assert os.path.exists(A_						)
       os.makedirs(A_,    exist_ok=A_						)
       print(F'''Writing results to {pytorch_dump_folder_path}'''						)

       # handle various types of models

       _lowerCamelCase						:    Optional[int]         =			basename(A_						)
       _lowerCamelCase						:    int         =			dirname(A_						)

       _lowerCamelCase						:    Any         =			fairseq.model_parallel.models.transformer.ModelParallelTransformerModel
       _lowerCamelCase						:    Any         =			cls.hub_models()
       _lowerCamelCase						:    Optional[Any]         =			{'''bpe''': '''fastbpe''', '''tokenizer''': '''moses'''}
       _lowerCamelCase						:    int         =			'''.'''
       # note: since the model dump is old, fairseq has upgraded its model some
       # time later, and it does a whole lot of rewrites and splits on the saved
       # weights, therefore we can't use torch.load() directly on the model file.
       # see: upgrade_state_dict(state_dict) in fairseq_model.py
       print(F'''using checkpoint {checkpoint_file}'''						)
       _lowerCamelCase						:    Dict         =			hub_utils.from_pretrained(
           A_,    A_,    A_,    archive_map=A_,    **A_						)

       _lowerCamelCase						:    Optional[Any]         =			vars(chkpt['''args''']['''model''']						)

       _lowerCamelCase						:    List[Any]         =			args['''source_lang''']
       _lowerCamelCase						:    Dict         =			args['''target_lang''']

       _lowerCamelCase						:    int         =			dirname(A_						)
       _lowerCamelCase						:    str         =			basename(A_						)

       # dicts
       _lowerCamelCase						:    Union[str, Any]         =			os.path.join(A_,    F'''dict.{src_lang}.txt'''						)
       _lowerCamelCase						:    Dict         =			os.path.join(A_,    F'''dict.{tgt_lang}.txt'''						)

       _lowerCamelCase						:    Union[str, Any]         =			Dictionary.load(A_						)
       _lowerCamelCase						:    Any         =			rewrite_dict_keys(src_dict.indices						)
       _lowerCamelCase						:    int         =			len(A_						)
       _lowerCamelCase						:    Optional[int]         =			os.path.join(A_,    '''vocab-src.json'''						)
       print(F'''Generating {src_vocab_file} of {src_vocab_size} of {src_lang} records'''						)
       with open(A_,    '''w''',    encoding='''utf-8'''						) as f:
              f.write(json.dumps(A_,    ensure_ascii=A_,    indent=A_						)						)

       # detect whether this is a do_lower_case situation, which can be derived by checking whether we
       # have at least one uppercase letter in the source vocab
       _lowerCamelCase						:    List[str]         =			True
       for k in src_vocab.keys():
              if not k.islower():
                     _lowerCamelCase						:    Tuple         =			False
                     break

       _lowerCamelCase						:    Optional[int]         =			Dictionary.load(A_						)
       _lowerCamelCase						:    Optional[Any]         =			rewrite_dict_keys(tgt_dict.indices						)
       _lowerCamelCase						:    Optional[Any]         =			len(A_						)
       _lowerCamelCase						:    int         =			os.path.join(A_,    '''vocab-tgt.json'''						)
       print(F'''Generating {tgt_vocab_file} of {tgt_vocab_size} of {tgt_lang} records'''						)
       with open(A_,    '''w''',    encoding='''utf-8'''						) as f:
              f.write(json.dumps(A_,    ensure_ascii=A_,    indent=A_						)						)

       # merges_file (bpecodes)
       _lowerCamelCase						:    Optional[Any]         =			os.path.join(A_,    VOCAB_FILES_NAMES['''merges_file''']						)
       for fn in ["bpecodes", "code"]:  # older fairseq called the merges file "code"
              _lowerCamelCase						:    List[str]         =			os.path.join(A_,    A_						)
              if os.path.exists(A_						):
                     break
       with open(A_,    encoding='''utf-8'''						) as fin:
              _lowerCamelCase						:    Union[str, Any]         =			fin.read()
       _lowerCamelCase						:    Optional[Any]         =			re.sub(R''' \d+$''',    '''''',    A_,    0,    re.M						)  # remove frequency number
       print(F'''Generating {merges_file}'''						)
       with open(A_,    '''w''',    encoding='''utf-8'''						) as fout:
              fout.write(A_						)

       # model config
       _lowerCamelCase						:    str         =			os.path.join(A_,    '''config.json'''						)

       # validate bpe/tokenizer config, as currently it's hardcoded to moses+fastbpe -
       # may have to modify the tokenizer if a different type is used by a future model
       assert args["bpe"] == "fastbpe", F'''need to extend tokenizer to support bpe={args["bpe"]}'''
       assert args["tokenizer"] == "moses", F'''need to extend tokenizer to support bpe={args["tokenizer"]}'''

       _lowerCamelCase						:    Optional[int]         =			{
           '''architectures''': ['''FSMTForConditionalGeneration'''],
           '''model_type''': '''fsmt''',
           '''activation_dropout''': args['''activation_dropout'''],
           '''activation_function''': '''relu''',
           '''attention_dropout''': args['''attention_dropout'''],
           '''d_model''': args['''decoder_embed_dim'''],
           '''dropout''': args['''dropout'''],
           '''init_std''': 0.02,
           '''max_position_embeddings''': args['''max_source_positions'''],
           '''num_hidden_layers''': args['''encoder_layers'''],
           '''src_vocab_size''': src_vocab_size,
           '''tgt_vocab_size''': tgt_vocab_size,
           '''langs''': [src_lang, tgt_lang],
           '''encoder_attention_heads''': args['''encoder_attention_heads'''],
           '''encoder_ffn_dim''': args['''encoder_ffn_embed_dim'''],
           '''encoder_layerdrop''': args['''encoder_layerdrop'''],
           '''encoder_layers''': args['''encoder_layers'''],
           '''decoder_attention_heads''': args['''decoder_attention_heads'''],
           '''decoder_ffn_dim''': args['''decoder_ffn_embed_dim'''],
           '''decoder_layerdrop''': args['''decoder_layerdrop'''],
           '''decoder_layers''': args['''decoder_layers'''],
           '''bos_token_id''': 0,
           '''pad_token_id''': 1,
           '''eos_token_id''': 2,
           '''is_encoder_decoder''': True,
           '''scale_embedding''': not args['''no_scale_embedding'''],
           '''tie_word_embeddings''': args['''share_all_embeddings'''],
       }

       # good hparam defaults to start with
       _lowerCamelCase						:    Tuple         =			5
       _lowerCamelCase						:    Dict         =			False
       if model_dir in best_score_hparams and "length_penalty" in best_score_hparams[model_dir]:
              _lowerCamelCase						:    Any         =			best_score_hparams[model_dir]['''length_penalty''']
       else:
              _lowerCamelCase						:    Optional[int]         =			1.0

       print(F'''Generating {fsmt_model_config_file}'''						)
       with open(A_,    '''w''',    encoding='''utf-8'''						) as f:
              f.write(json.dumps(A_,    ensure_ascii=A_,    indent=A_						)						)

       # tokenizer config
       _lowerCamelCase						:    Dict         =			os.path.join(A_,    A_						)

       _lowerCamelCase						:    List[str]         =			{
           '''langs''': [src_lang, tgt_lang],
           '''model_max_length''': 10_24,
           '''do_lower_case''': do_lower_case,
       }

       print(F'''Generating {fsmt_tokenizer_config_file}'''						)
       with open(A_,    '''w''',    encoding='''utf-8'''						) as f:
              f.write(json.dumps(A_,    ensure_ascii=A_,    indent=A_						)						)

       # model
       _lowerCamelCase						:    str         =			chkpt['''models'''][0]
       _lowerCamelCase						:    List[Any]         =			model.state_dict()

       # rename keys to start with 'model.'
       _lowerCamelCase						:    str         =			OrderedDict(('''model.''' + k, v) for k, v in model_state_dict.items()						)

       # remove unneeded keys
       _lowerCamelCase						:    str         =			[
           '''model.model''',
           '''model.encoder.version''',
           '''model.decoder.version''',
           '''model.encoder_embed_tokens.weight''',
           '''model.decoder_embed_tokens.weight''',
           '''model.encoder.embed_positions._float_tensor''',
           '''model.decoder.embed_positions._float_tensor''',
       ]
       for k in ignore_keys:
              model_state_dict.pop(A_,    A_						)

       _lowerCamelCase						:    int         =			FSMTConfig.from_pretrained(A_						)
       _lowerCamelCase						:    List[str]         =			FSMTForConditionalGeneration(A_						)

       # check that it loads ok
       model_new.load_state_dict(A_,    strict=A_						)

       # save
       _lowerCamelCase						:    Any         =			os.path.join(A_,    A_						)
       print(F'''Generating {pytorch_weights_dump_path}'''						)
       torch.save(A_,    A_						)

       print('''Conversion is done!'''						)
       print('''\nLast step is to upload the files to s3'''						)
       print(F'''cd {data_root}'''						)
       print(F'''transformers-cli upload {model_dir}'''						)


if __name__ == "__main__":
   lowerCAmelCase__      =    argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '''--fsmt_checkpoint_path''',
       default=None,
       type=str,
       required=True,
       help=(
           '''Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,'''
           ''' bpecodes, etc.'''
       ),
   )
   parser.add_argument(
       '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
   )
   lowerCAmelCase__      =    parser.parse_args()
   convert_fsmt_checkpoint_to_pytorch(args.fsmt_checkpoint_path, args.pytorch_dump_folder_path)





"""simple docstring"""






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 	__snake_case	(     _lowercase  ,					unittest.TestCase):
  # TODO: is there an appropriate internal test set?
  snake_case__    :       List[str]        =						"ssube/stable-diffusion-x4-upscaler-onnx"




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						int=0		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			floats_tensor((1, 3, 1_2_8, 1_2_8)						, rng=random.Random(__lowerCAmelCase		)		)
         _lowerCamelCase						:    Union[str, Any]         =			torch.manual_seed(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			{
             '''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 						SCREAMING_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						:    Union[str, Any]         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Any         =			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_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 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






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

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Optional[int]         =			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 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






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

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			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 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    Union[str, Any]         =			EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Union[str, 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 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    int         =			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Optional[int]         =			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 	__snake_case	(     unittest.TestCase):




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






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




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			ort.SessionOptions()
         _lowerCamelCase						:    List[str]         =			False
         return options




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






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

         _lowerCamelCase						:    List[Any]         =			torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			pipe(
             prompt=__lowerCAmelCase						, image=__lowerCAmelCase						, guidance_scale=7.5						, num_inference_steps=1_0						, generator=__lowerCAmelCase						, output_type='''np'''						, )
         _lowerCamelCase						:    List[Any]         =			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						:    str         =			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 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''		)
         _lowerCamelCase						:    int         =			init_image.resize((1_2_8, 1_2_8)		)
         _lowerCamelCase						:    str         =			LMSDiscreteScheduler.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''						, subfolder='''scheduler'''		)
         _lowerCamelCase						:    Dict         =			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						:    Optional[int]         =			'''A fantasy landscape, trending on artstation'''

         _lowerCamelCase						:    int         =			torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			pipe(
             prompt=__lowerCAmelCase						, image=__lowerCAmelCase						, guidance_scale=7.5						, num_inference_steps=2_0						, generator=__lowerCAmelCase						, output_type='''np'''						, )
         _lowerCamelCase						:    Union[str, Any]         =			output.images
         _lowerCamelCase						:    Optional[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						:    str         =			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


