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    (						lowerCamelCase_       ,		unittest.TestCase		):
   # TODO: is there an appropriate internal test set?
   lowerCAmelCase_														=							"ssube/stable-diffusion-x4-upscaler-onnx"




   def      __a						(   self      :    Dict	,     _lowercase      :    Any=0				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				floats_tensor((1, 3, 1_28, 1_28)	,     rng=random.Random(_lowercase				)				)
     SCREAMING_SNAKE_CASE__      							=				torch.manual_seed(_lowercase				)
     SCREAMING_SNAKE_CASE__      							=				{
         """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[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint	,     provider="""CPUExecutionProvider"""				)
     pipe.set_progress_bar_config(disable=_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE__      							=				pipe(**_lowercase				).images
     SCREAMING_SNAKE_CASE__      							=				image[0, -3:, -3:, -1].flatten()

     # started as 128, should now be 512
     assert image.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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      :    List[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint	,     provider="""CPUExecutionProvider"""				)
     SCREAMING_SNAKE_CASE__      							=				PNDMScheduler.from_config(pipe.scheduler.config	,     skip_prk_steps=_lowercase				)
     pipe.set_progress_bar_config(disable=_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE__      							=				pipe(**_lowercase				).images
     SCREAMING_SNAKE_CASE__      							=				image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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[int]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint	,     provider="""CPUExecutionProvider"""				)
     SCREAMING_SNAKE_CASE__      							=				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
     pipe.set_progress_bar_config(disable=_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE__      							=				pipe(**_lowercase				).images
     SCREAMING_SNAKE_CASE__      							=				image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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      :    List[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint	,     provider="""CPUExecutionProvider"""				)
     SCREAMING_SNAKE_CASE__      							=				EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
     pipe.set_progress_bar_config(disable=_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE__      							=				pipe(**_lowercase				).images
     SCREAMING_SNAKE_CASE__      							=				image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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      :    Any				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint	,     provider="""CPUExecutionProvider"""				)
     SCREAMING_SNAKE_CASE__      							=				EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config				)
     pipe.set_progress_bar_config(disable=_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE__      							=				pipe(**_lowercase				).images
     SCREAMING_SNAKE_CASE__      							=				image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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      __a						(   self      :    Dict				):



     """simple docstring"""






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




   @property
   def      __a						(   self      :    Any				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				ort.SessionOptions()
     SCREAMING_SNAKE_CASE__      							=				False
     return options




   def      __a						(   self      :    Optional[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				load_image(
         """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
         """/img2img/sketch-mountains-input.jpg"""				)
     SCREAMING_SNAKE_CASE__      							=				init_image.resize((1_28, 1_28)				)
     # using the PNDM scheduler by default
     SCREAMING_SNAKE_CASE__      							=				OnnxStableDiffusionUpscalePipeline.from_pretrained(
         """ssube/stable-diffusion-x4-upscaler-onnx"""	,     provider=self.gpu_provider	,     sess_options=self.gpu_options	,     )
     pipe.set_progress_bar_config(disable=_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				"""A fantasy landscape, trending on artstation"""

     SCREAMING_SNAKE_CASE__      							=				torch.manual_seed(0				)
     SCREAMING_SNAKE_CASE__      							=				pipe(
         prompt=_lowercase	,     image=_lowercase	,     guidance_scale=7.5	,     num_inference_steps=10	,     generator=_lowercase	,     output_type="""np"""	,     )
     SCREAMING_SNAKE_CASE__      							=				output.images
     SCREAMING_SNAKE_CASE__      							=				images[0, 2_55:2_58, 3_83:3_86, -1]

     assert images.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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      :    List[Any]				):



     """simple docstring"""






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

     SCREAMING_SNAKE_CASE__      							=				"""A fantasy landscape, trending on artstation"""

     SCREAMING_SNAKE_CASE__      							=				torch.manual_seed(0				)
     SCREAMING_SNAKE_CASE__      							=				pipe(
         prompt=_lowercase	,     image=_lowercase	,     guidance_scale=7.5	,     num_inference_steps=20	,     generator=_lowercase	,     output_type="""np"""	,     )
     SCREAMING_SNAKE_CASE__      							=				output.images
     SCREAMING_SNAKE_CASE__      							=				images[0, 2_55:2_58, 3_83:3_86, -1]

     assert images.shape == (1, 5_12, 5_12, 3)
     SCREAMING_SNAKE_CASE__      							=				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 io

import requests
import torch
from omegaconf import OmegaConf

from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    assign_to_checkpoint,
    conv_attn_to_linear,
    create_vae_diffusers_config,
    renew_vae_attention_paths,
    renew_vae_resnet_paths,
)



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Optional[int]			,					__UpperCamelCase   :	str     )							-> Tuple:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				checkpoint

  SCREAMING_SNAKE_CASE__      							=				{}

  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_in.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_in.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_out.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.norm_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.norm_out.bias"""]

  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_in.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_in.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_out.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.norm_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.norm_out.bias"""]

  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""quant_conv.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""quant_conv.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""post_quant_conv.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""post_quant_conv.bias"""]

  # Retrieves the keys for the encoder down blocks only
  SCREAMING_SNAKE_CASE__      							=				len({""".""".join(layer.split("""."""     )[:3]     ) for layer in vae_state_dict if """encoder.down""" in layer}     )
  SCREAMING_SNAKE_CASE__      							=				{
      layer_id: [key for key in vae_state_dict if f"""down.{layer_id}""" in key] for layer_id in range(__UpperCamelCase     )
  }

  # Retrieves the keys for the decoder up blocks only
  SCREAMING_SNAKE_CASE__      							=				len({""".""".join(layer.split("""."""     )[:3]     ) for layer in vae_state_dict if """decoder.up""" in layer}     )
  SCREAMING_SNAKE_CASE__      							=				{
      layer_id: [key for key in vae_state_dict if f"""up.{layer_id}""" in key] for layer_id in range(__UpperCamelCase     )
  }

  for i in range(__UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__      							=				[key for key in down_blocks[i] if f"""down.{i}""" in key and f"""down.{i}.downsample""" not in key]

    if f"""encoder.down.{i}.downsample.conv.weight""" in vae_state_dict:
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict.pop(
          f"""encoder.down.{i}.downsample.conv.weight"""     )
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict.pop(
          f"""encoder.down.{i}.downsample.conv.bias"""     )

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""down.{i}.block""", """new""": f"""down_blocks.{i}.resnets"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """encoder.mid.block""" in key]
  SCREAMING_SNAKE_CASE__      							=				2
  for i in range(1			,					num_mid_res_blocks + 1     ):
    SCREAMING_SNAKE_CASE__      							=				[key for key in mid_resnets if f"""encoder.mid.block_{i}""" in key]

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""mid.block_{i}""", """new""": f"""mid_block.resnets.{i - 1}"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """encoder.mid.attn""" in key]
  SCREAMING_SNAKE_CASE__      							=				renew_vae_attention_paths(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				{"""old""": """mid.attn_1""", """new""": """mid_block.attentions.0"""}
  assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )
  conv_attn_to_linear(__UpperCamelCase     )

  for i in range(__UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__      							=				num_up_blocks - 1 - i
    SCREAMING_SNAKE_CASE__      							=				[
        key for key in up_blocks[block_id] if f"""up.{block_id}""" in key and f"""up.{block_id}.upsample""" not in key
    ]

    if f"""decoder.up.{block_id}.upsample.conv.weight""" in vae_state_dict:
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict[
          f"""decoder.up.{block_id}.upsample.conv.weight"""
      ]
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict[
          f"""decoder.up.{block_id}.upsample.conv.bias"""
      ]

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""up.{block_id}.block""", """new""": f"""up_blocks.{i}.resnets"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """decoder.mid.block""" in key]
  SCREAMING_SNAKE_CASE__      							=				2
  for i in range(1			,					num_mid_res_blocks + 1     ):
    SCREAMING_SNAKE_CASE__      							=				[key for key in mid_resnets if f"""decoder.mid.block_{i}""" in key]

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""mid.block_{i}""", """new""": f"""mid_block.resnets.{i - 1}"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """decoder.mid.attn""" in key]
  SCREAMING_SNAKE_CASE__      							=				renew_vae_attention_paths(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				{"""old""": """mid.attn_1""", """new""": """mid_block.attentions.0"""}
  assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )
  conv_attn_to_linear(__UpperCamelCase     )
  return new_checkpoint



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	str			,					__UpperCamelCase   :	str			,					)							-> Optional[int]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				requests.get(
      """ https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"""     )
  SCREAMING_SNAKE_CASE__      							=				io.BytesIO(r.content     )

  SCREAMING_SNAKE_CASE__      							=				OmegaConf.load(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				5_12
  SCREAMING_SNAKE_CASE__      							=				"""cuda""" if torch.cuda.is_available() else """cpu"""
  if checkpoint_path.endswith("""safetensors"""     ):
    from safetensors import safe_open

    SCREAMING_SNAKE_CASE__      							=				{}
    with safe_open(__UpperCamelCase			,					framework="""pt"""			,					device="""cpu"""     ) as f:
      for key in f.keys():
        SCREAMING_SNAKE_CASE__      							=				f.get_tensor(__UpperCamelCase     )
  else:
    SCREAMING_SNAKE_CASE__      							=				torch.load(__UpperCamelCase			,					map_location=__UpperCamelCase     )["""state_dict"""]

  # Convert the VAE model.
  SCREAMING_SNAKE_CASE__      							=				create_vae_diffusers_config(__UpperCamelCase			,					image_size=__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				custom_convert_ldm_vae_checkpoint(__UpperCamelCase			,					__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				AutoencoderKL(**__UpperCamelCase     )
  vae.load_state_dict(__UpperCamelCase     )
  vae.save_pretrained(__UpperCamelCase     )


if __name__ == "__main__":
   __lowerCamelCase      :			List[Any]							       =		argparse.ArgumentParser()

   parser.add_argument('''--vae_pt_path''', default=None, type=str, required=True, help='''Path to the VAE.pt to convert.''')
   parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the VAE.pt to convert.''')

   __lowerCamelCase      :			Optional[int]							       =		parser.parse_args()

   vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)
