'''simple docstring'''




import argparse
import os
import shutil
from pathlib import Path

import onnx
import torch
from packaging import version
from torch.onnx import export

from diffusers import OnnxRuntimeModel, OnnxStableDiffusionPipeline, StableDiffusionPipeline


__UpperCAmelCase												=version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__=False    ,  )     ->   List[Any]:
    output_path.parent.mkdir(parents=UpperCamelCase__    ,  exist_ok=UpperCamelCase__    )
    # PyTorch deprecated the `enable_onnx_checker` and `use_external_data_format` arguments in v1.11,
    # so we check the torch version for backwards compatibility
    if is_torch_less_than_1_11:
        export(
            UpperCamelCase__    ,  UpperCamelCase__    ,  f=output_path.as_posix()    ,  input_names=UpperCamelCase__    ,  output_names=UpperCamelCase__    ,  dynamic_axes=UpperCamelCase__    ,  do_constant_folding=UpperCamelCase__    ,  use_external_data_format=UpperCamelCase__    ,  enable_onnx_checker=UpperCamelCase__    ,  opset_version=UpperCamelCase__    ,  )
    else:
        export(
            UpperCamelCase__    ,  UpperCamelCase__    ,  f=output_path.as_posix()    ,  input_names=UpperCamelCase__    ,  output_names=UpperCamelCase__    ,  dynamic_axes=UpperCamelCase__    ,  do_constant_folding=UpperCamelCase__    ,  opset_version=UpperCamelCase__    ,  )






@torch.no_grad()
def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__ = False    )     ->   Tuple:
    __lowerCamelCase          =  torch.floataa if fpaa else torch.floataa
    if fpaa and torch.cuda.is_available():
        __lowerCamelCase          =  '''cuda'''
    elif fpaa and not torch.cuda.is_available():
        raise ValueError('''`float16` model export is only supported on GPUs with CUDA'''    )
    else:
        __lowerCamelCase          =  '''cpu'''
    __lowerCamelCase          =  StableDiffusionPipeline.from_pretrained(UpperCamelCase__    ,  torch_dtype=UpperCamelCase__    ).to(UpperCamelCase__    )
    __lowerCamelCase          =  Path(UpperCamelCase__    )

    # TEXT ENCODER
    __lowerCamelCase          =  pipeline.text_encoder.config.max_position_embeddings
    __lowerCamelCase          =  pipeline.text_encoder.config.hidden_size
    __lowerCamelCase          =  pipeline.tokenizer(
        '''A sample prompt'''    ,  padding='''max_length'''    ,  max_length=pipeline.tokenizer.model_max_length    ,  truncation=UpperCamelCase__    ,  return_tensors='''pt'''    ,  )
    onnx_export(
        pipeline.text_encoder    ,  model_args=(text_input.input_ids.to(device=UpperCamelCase__    ,  dtype=torch.intaa    ))    ,  output_path=output_path / '''text_encoder''' / '''model.onnx'''    ,  ordered_input_names=['''input_ids''']    ,  output_names=['''last_hidden_state''', '''pooler_output''']    ,  dynamic_axes={
            '''input_ids''': {0: '''batch''', 1: '''sequence'''},
        }    ,  opset=UpperCamelCase__    ,  )
    del pipeline.text_encoder

    # UNET
    __lowerCamelCase          =  pipeline.unet.config.in_channels
    __lowerCamelCase          =  pipeline.unet.config.sample_size
    __lowerCamelCase          =  output_path / '''unet''' / '''model.onnx'''
    onnx_export(
        pipeline.unet    ,  model_args=(
            torch.randn(2    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
            torch.randn(2    ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
            torch.randn(2    ,  UpperCamelCase__    ,  UpperCamelCase__    ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
            False,
        )    ,  output_path=UpperCamelCase__    ,  ordered_input_names=['''sample''', '''timestep''', '''encoder_hidden_states''', '''return_dict''']    ,  output_names=['''out_sample''']    ,  dynamic_axes={
            '''sample''': {0: '''batch''', 1: '''channels''', 2: '''height''', 3: '''width'''},
            '''timestep''': {0: '''batch'''},
            '''encoder_hidden_states''': {0: '''batch''', 1: '''sequence'''},
        }    ,  opset=UpperCamelCase__    ,  use_external_data_format=UpperCamelCase__    ,  )
    __lowerCamelCase          =  str(unet_path.absolute().as_posix()    )
    __lowerCamelCase          =  os.path.dirname(UpperCamelCase__    )
    __lowerCamelCase          =  onnx.load(UpperCamelCase__    )
    # clean up existing tensor files
    shutil.rmtree(UpperCamelCase__    )
    os.mkdir(UpperCamelCase__    )
    # collate external tensor files into one
    onnx.save_model(
        UpperCamelCase__    ,  UpperCamelCase__    ,  save_as_external_data=UpperCamelCase__    ,  all_tensors_to_one_file=UpperCamelCase__    ,  location='''weights.pb'''    ,  convert_attribute=UpperCamelCase__    ,  )
    del pipeline.unet

    # VAE ENCODER
    __lowerCamelCase          =  pipeline.vae
    __lowerCamelCase          =  vae_encoder.config.in_channels
    __lowerCamelCase          =  vae_encoder.config.sample_size
    # need to get the raw tensor output (sample) from the encoder
    __lowerCamelCase          =  lambda UpperCamelCase__    ,  UpperCamelCase__    : vae_encoder.encode(UpperCamelCase__    ,  UpperCamelCase__    )[0].sample()
    onnx_export(
        UpperCamelCase__    ,  model_args=(
            torch.randn(1    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
            False,
        )    ,  output_path=output_path / '''vae_encoder''' / '''model.onnx'''    ,  ordered_input_names=['''sample''', '''return_dict''']    ,  output_names=['''latent_sample''']    ,  dynamic_axes={
            '''sample''': {0: '''batch''', 1: '''channels''', 2: '''height''', 3: '''width'''},
        }    ,  opset=UpperCamelCase__    ,  )

    # VAE DECODER
    __lowerCamelCase          =  pipeline.vae
    __lowerCamelCase          =  vae_decoder.config.latent_channels
    __lowerCamelCase          =  vae_decoder.config.out_channels
    # forward only through the decoder part
    __lowerCamelCase          =  vae_encoder.decode
    onnx_export(
        UpperCamelCase__    ,  model_args=(
            torch.randn(1    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
            False,
        )    ,  output_path=output_path / '''vae_decoder''' / '''model.onnx'''    ,  ordered_input_names=['''latent_sample''', '''return_dict''']    ,  output_names=['''sample''']    ,  dynamic_axes={
            '''latent_sample''': {0: '''batch''', 1: '''channels''', 2: '''height''', 3: '''width'''},
        }    ,  opset=UpperCamelCase__    ,  )
    del pipeline.vae

    # SAFETY CHECKER
    if pipeline.safety_checker is not None:
        __lowerCamelCase          =  pipeline.safety_checker
        __lowerCamelCase          =  safety_checker.config.vision_config.num_channels
        __lowerCamelCase          =  safety_checker.config.vision_config.image_size
        __lowerCamelCase          =  safety_checker.forward_onnx
        onnx_export(
            pipeline.safety_checker    ,  model_args=(
                torch.randn(
                    1    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
                torch.randn(1    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ).to(device=UpperCamelCase__    ,  dtype=UpperCamelCase__    ),
            )    ,  output_path=output_path / '''safety_checker''' / '''model.onnx'''    ,  ordered_input_names=['''clip_input''', '''images''']    ,  output_names=['''out_images''', '''has_nsfw_concepts''']    ,  dynamic_axes={
                '''clip_input''': {0: '''batch''', 1: '''channels''', 2: '''height''', 3: '''width'''},
                '''images''': {0: '''batch''', 1: '''height''', 2: '''width''', 3: '''channels'''},
            }    ,  opset=UpperCamelCase__    ,  )
        del pipeline.safety_checker
        __lowerCamelCase          =  OnnxRuntimeModel.from_pretrained(output_path / '''safety_checker'''    )
        __lowerCamelCase          =  pipeline.feature_extractor
    else:
        __lowerCamelCase          =  None
        __lowerCamelCase          =  None

    __lowerCamelCase          =  OnnxStableDiffusionPipeline(
        vae_encoder=OnnxRuntimeModel.from_pretrained(output_path / '''vae_encoder'''    )    ,  vae_decoder=OnnxRuntimeModel.from_pretrained(output_path / '''vae_decoder'''    )    ,  text_encoder=OnnxRuntimeModel.from_pretrained(output_path / '''text_encoder'''    )    ,  tokenizer=pipeline.tokenizer    ,  unet=OnnxRuntimeModel.from_pretrained(output_path / '''unet'''    )    ,  scheduler=pipeline.scheduler    ,  safety_checker=UpperCamelCase__    ,  feature_extractor=UpperCamelCase__    ,  requires_safety_checker=safety_checker is not None    ,  )

    onnx_pipeline.save_pretrained(UpperCamelCase__    )
    print('''ONNX pipeline saved to'''    ,  UpperCamelCase__    )

    del pipeline
    del onnx_pipeline
    __lowerCamelCase          =  OnnxStableDiffusionPipeline.from_pretrained(UpperCamelCase__    ,  provider='''CPUExecutionProvider'''    )
    print('''ONNX pipeline is loadable'''    )


if __name__ == "__main__":
  __UpperCAmelCase												=argparse.ArgumentParser()

  parser.add_argument(
      "--model_path",
      type=str,
      required=True,
      help="Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).",
  )

  parser.add_argument("--output_path", type=str, required=True, help="Path to the output model.")

  parser.add_argument(
      "--opset",
      default=1_4,
      type=int,
      help="The version of the ONNX operator set to use.",
  )
  parser.add_argument("--fp16", action="store_true", default=False, help="Export the models in `float16` mode")

  __UpperCAmelCase												=parser.parse_args()

  convert_models(args.model_path, args.output_path, args.opset, args.fpaa)

'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase												={
    "configuration_whisper": ["WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP", "WhisperConfig", "WhisperOnnxConfig"],
    "feature_extraction_whisper": ["WhisperFeatureExtractor"],
    "processing_whisper": ["WhisperProcessor"],
    "tokenization_whisper": ["WhisperTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=["WhisperTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "WhisperForConditionalGeneration",
      "WhisperModel",
      "WhisperPreTrainedModel",
      "WhisperForAudioClassification",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFWhisperForConditionalGeneration",
      "TFWhisperModel",
      "TFWhisperPreTrainedModel",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "FlaxWhisperForConditionalGeneration",
      "FlaxWhisperModel",
      "FlaxWhisperPreTrainedModel",
      "FlaxWhisperForAudioClassification",
  ]


if TYPE_CHECKING:
  from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
  from .feature_extraction_whisper import WhisperFeatureExtractor
  from .processing_whisper import WhisperProcessor
  from .tokenization_whisper import WhisperTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_whisper_fast import WhisperTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_whisper import (
        WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
        WhisperForAudioClassification,
        WhisperForConditionalGeneration,
        WhisperModel,
        WhisperPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_whisper import (
        TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFWhisperForConditionalGeneration,
        TFWhisperModel,
        TFWhisperPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_whisper import (
        FlaxWhisperForAudioClassification,
        FlaxWhisperForConditionalGeneration,
        FlaxWhisperModel,
        FlaxWhisperPreTrainedModel,
    )

else:
  import sys

  __UpperCAmelCase												=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
