'''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


SCREAMING_SNAKE_CASE_	   =			version.parse(version.parse(torch.__version__).base_version) < version.parse('1.11')






def  UpperCamelCase__					(					_lowercase    :				int   ,    _lowercase    :				tuple   ,    _lowercase    :				Path   ,    _lowercase    :				List[Any]   ,    _lowercase    :				List[str]   ,    _lowercase    :				List[str]   ,    _lowercase    :				Dict   ,    _lowercase    :				Any=False   ,    )							->	Dict:
  output_path.parent.mkdir(parents=_lowercase   ,    exist_ok=_lowercase     )
  # 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(
        _lowercase   ,    _lowercase   ,    f=output_path.as_posix()   ,    input_names=_lowercase   ,    output_names=_lowercase   ,    dynamic_axes=_lowercase   ,    do_constant_folding=_lowercase   ,    use_external_data_format=_lowercase   ,    enable_onnx_checker=_lowercase   ,    opset_version=_lowercase   ,    )
  else:
    export(
        _lowercase   ,    _lowercase   ,    f=output_path.as_posix()   ,    input_names=_lowercase   ,    output_names=_lowercase   ,    dynamic_axes=_lowercase   ,    do_constant_folding=_lowercase   ,    opset_version=_lowercase   ,    )






@torch.no_grad()
def  UpperCamelCase__					(					_lowercase    :				str   ,    _lowercase    :				str   ,    _lowercase    :				int   ,    _lowercase    :				bool = False     )							->	int:
  __UpperCAmelCase:					Optional[Any]		= torch.floataa if fpaa else torch.floataa
  if fpaa and torch.cuda.is_available():
    __UpperCAmelCase:					Union[str, Any]		= 'cuda'
  elif fpaa and not torch.cuda.is_available():
    raise ValueError("""`float16` model export is only supported on GPUs with CUDA"""     )
  else:
    __UpperCAmelCase:					int		= 'cpu'
  __UpperCAmelCase:					Any		= StableDiffusionPipeline.from_pretrained(_lowercase   ,    torch_dtype=_lowercase     ).to(_lowercase     )
  __UpperCAmelCase:					Optional[int]		= Path(_lowercase     )

  # TEXT ENCODER
  __UpperCAmelCase:					Any		= pipeline.text_encoder.config.max_position_embeddings
  __UpperCAmelCase:					Tuple		= pipeline.text_encoder.config.hidden_size
  __UpperCAmelCase:					str		= pipeline.tokenizer(
      """A sample prompt"""   ,    padding="""max_length"""   ,    max_length=pipeline.tokenizer.model_max_length   ,    truncation=_lowercase   ,    return_tensors="""pt"""   ,    )
  onnx_export(
      pipeline.text_encoder   ,    model_args=(text_input.input_ids.to(device=_lowercase   ,    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=_lowercase   ,    )
  del pipeline.text_encoder

  # UNET
  __UpperCAmelCase:					List[str]		= pipeline.unet.config.in_channels
  __UpperCAmelCase:					str		= pipeline.unet.config.sample_size
  __UpperCAmelCase:					List[str]		= output_path / 'unet' / 'model.onnx'
  onnx_export(
      pipeline.unet   ,    model_args=(
          torch.randn(2   ,    _lowercase   ,    _lowercase   ,    _lowercase     ).to(device=_lowercase   ,    dtype=_lowercase     ),
          torch.randn(2     ).to(device=_lowercase   ,    dtype=_lowercase     ),
          torch.randn(2   ,    _lowercase   ,    _lowercase     ).to(device=_lowercase   ,    dtype=_lowercase     ),
          False,
      )   ,    output_path=_lowercase   ,    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=_lowercase   ,    use_external_data_format=_lowercase   ,    )
  __UpperCAmelCase:					Optional[int]		= str(unet_path.absolute().as_posix()     )
  __UpperCAmelCase:					List[Any]		= os.path.dirname(_lowercase     )
  __UpperCAmelCase:					List[str]		= onnx.load(_lowercase     )
  # clean up existing tensor files
  shutil.rmtree(_lowercase     )
  os.mkdir(_lowercase     )
  # collate external tensor files into one
  onnx.save_model(
      _lowercase   ,    _lowercase   ,    save_as_external_data=_lowercase   ,    all_tensors_to_one_file=_lowercase   ,    location="""weights.pb"""   ,    convert_attribute=_lowercase   ,    )
  del pipeline.unet

  # VAE ENCODER
  __UpperCAmelCase:					Any		= pipeline.vae
  __UpperCAmelCase:					Optional[Any]		= vae_encoder.config.in_channels
  __UpperCAmelCase:					List[Any]		= vae_encoder.config.sample_size
  # need to get the raw tensor output (sample) from the encoder
  __UpperCAmelCase:					Dict		= lambda _lowercase   ,    _lowercase     : vae_encoder.encode(_lowercase   ,    _lowercase     )[0].sample()
  onnx_export(
      _lowercase   ,    model_args=(
          torch.randn(1   ,    _lowercase   ,    _lowercase   ,    _lowercase     ).to(device=_lowercase   ,    dtype=_lowercase     ),
          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=_lowercase   ,    )

  # VAE DECODER
  __UpperCAmelCase:					List[str]		= pipeline.vae
  __UpperCAmelCase:					Any		= vae_decoder.config.latent_channels
  __UpperCAmelCase:					Optional[Any]		= vae_decoder.config.out_channels
  # forward only through the decoder part
  __UpperCAmelCase:					List[str]		= vae_encoder.decode
  onnx_export(
      _lowercase   ,    model_args=(
          torch.randn(1   ,    _lowercase   ,    _lowercase   ,    _lowercase     ).to(device=_lowercase   ,    dtype=_lowercase     ),
          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=_lowercase   ,    )
  del pipeline.vae

  # SAFETY CHECKER
  if pipeline.safety_checker is not None:
    __UpperCAmelCase:					int		= pipeline.safety_checker
    __UpperCAmelCase:					List[str]		= safety_checker.config.vision_config.num_channels
    __UpperCAmelCase:					str		= safety_checker.config.vision_config.image_size
    __UpperCAmelCase:					Tuple		= safety_checker.forward_onnx
    onnx_export(
        pipeline.safety_checker   ,    model_args=(
            torch.randn(
                1   ,    _lowercase   ,    _lowercase   ,    _lowercase   ,    ).to(device=_lowercase   ,    dtype=_lowercase     ),
            torch.randn(1   ,    _lowercase   ,    _lowercase   ,    _lowercase     ).to(device=_lowercase   ,    dtype=_lowercase     ),
        )   ,    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=_lowercase   ,    )
    del pipeline.safety_checker
    __UpperCAmelCase:					Any		= OnnxRuntimeModel.from_pretrained(output_path / """safety_checker"""     )
    __UpperCAmelCase:					Optional[Any]		= pipeline.feature_extractor
  else:
    __UpperCAmelCase:					List[str]		= None
    __UpperCAmelCase:					List[str]		= None

  __UpperCAmelCase:					Dict		= 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=_lowercase   ,    feature_extractor=_lowercase   ,    requires_safety_checker=safety_checker is not None   ,    )

  onnx_pipeline.save_pretrained(_lowercase     )
  print("""ONNX pipeline saved to"""   ,    _lowercase     )

  del pipeline
  del onnx_pipeline
  __UpperCAmelCase:					Any		= OnnxStableDiffusionPipeline.from_pretrained(_lowercase   ,    provider="""CPUExecutionProvider"""     )
  print("""ONNX pipeline is loadable"""     )


if __name__ == "__main__":
  SCREAMING_SNAKE_CASE_	   =			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=14,
      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')

  SCREAMING_SNAKE_CASE_	   =			parser.parse_args()

  convert_models(args.model_path, args.output_path, args.opset, args.fpaa)
'''simple docstring'''






from __future__ import annotations

SCREAMING_SNAKE_CASE_	   =			10






def  UpperCamelCase__					(					_lowercase    :				list[int]     )							->	list[int]:
  __UpperCAmelCase:					Union[str, Any]		= 1
  __UpperCAmelCase:					Optional[Any]		= max(_lowercase     )
  while placement <= max_digit:
    # declare and initialize empty buckets
    __UpperCAmelCase:					list[list]		= [[] for _ in range(_lowercase     )]
    # split list_of_ints between the buckets
    for i in list_of_ints:
      __UpperCAmelCase:					Optional[Any]		= int((i / placement) % RADIX     )
      buckets[tmp].append(_lowercase     )
    # put each buckets' contents into list_of_ints
    __UpperCAmelCase:					Optional[int]		= 0
    for b in range(_lowercase     ):
      for i in buckets[b]:
        __UpperCAmelCase:					str		= i
        a += 1
        # move to next
    placement *= RADIX
  return list_of_ints


if __name__ == "__main__":
    import doctest

    doctest.testmod()