



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


_lowerCAmelCase							:   int				=				version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Tuple		,    SCREAMING_SNAKE_CASE__      :		tuple		,    SCREAMING_SNAKE_CASE__      :		Path		,    SCREAMING_SNAKE_CASE__      :		Any		,    SCREAMING_SNAKE_CASE__      :		List[str]		,    SCREAMING_SNAKE_CASE__      :		str		,    SCREAMING_SNAKE_CASE__      :		Dict		,    SCREAMING_SNAKE_CASE__      :		Tuple=False		,    )    ->  Tuple:





     '''simple docstring'''

     output_path.parent.mkdir(parents=SCREAMING_SNAKE_CASE__		,    exist_ok=SCREAMING_SNAKE_CASE__							)
     # 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(
              SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    f=output_path.as_posix()		,    input_names=SCREAMING_SNAKE_CASE__		,    output_names=SCREAMING_SNAKE_CASE__		,    dynamic_axes=SCREAMING_SNAKE_CASE__		,    do_constant_folding=SCREAMING_SNAKE_CASE__		,    use_external_data_format=SCREAMING_SNAKE_CASE__		,    enable_onnx_checker=SCREAMING_SNAKE_CASE__		,    opset_version=SCREAMING_SNAKE_CASE__		,    )
     else:
          export(
              SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    f=output_path.as_posix()		,    input_names=SCREAMING_SNAKE_CASE__		,    output_names=SCREAMING_SNAKE_CASE__		,    dynamic_axes=SCREAMING_SNAKE_CASE__		,    do_constant_folding=SCREAMING_SNAKE_CASE__		,    opset_version=SCREAMING_SNAKE_CASE__		,    )

@torch.no_grad()
def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		str		,    SCREAMING_SNAKE_CASE__      :		str		,    SCREAMING_SNAKE_CASE__      :		int		,    SCREAMING_SNAKE_CASE__      :		bool = False							)    ->  Tuple:





     '''simple docstring'''

     _UpperCAmelCase     :      Tuple														=       torch.floataa if fpaa else torch.floataa
     if fpaa and torch.cuda.is_available():
          _UpperCAmelCase     :      str														=       "cuda"
     elif fpaa and not torch.cuda.is_available():
          raise ValueError("`float16` model export is only supported on GPUs with CUDA"							)
     else:
          _UpperCAmelCase     :      Dict														=       "cpu"
     _UpperCAmelCase     :      Optional[Any]														=       StableDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE__		,    torch_dtype=SCREAMING_SNAKE_CASE__							).to(SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      Optional[int]														=       Path(SCREAMING_SNAKE_CASE__							)

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

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

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

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

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

     _UpperCAmelCase     :      Union[str, Any]														=       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=SCREAMING_SNAKE_CASE__		,    feature_extractor=SCREAMING_SNAKE_CASE__		,    requires_safety_checker=safety_checker is not None		,    )

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

     del pipeline
     del onnx_pipeline
     _UpperCAmelCase     :      Dict														=       OnnxStableDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE__		,    provider="CPUExecutionProvider"							)
     print("ONNX pipeline is loadable"							)


if __name__ == "__main__":
 _lowerCAmelCase							:   int				=				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")

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

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





"""simple docstring"""




import requests

_lowerCAmelCase							:   List[Any]				=				"https://newsapi.org/v1/articles?source=bbc-news&sortBy=top&apiKey="

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		str							)    ->  None:





     '''simple docstring'''

     _UpperCAmelCase     :      List[str]														=       requests.get(_NEWS_API + bbc_news_api_key							).json()
     # each article in the list is a dict
     for i, article in enumerate(bbc_news_page["articles"]		,    1							):
          print(f'{i}.) {article["title"]}'							)


if __name__ == "__main__":
 fetch_bbc_news(bbc_news_api_key="<Your BBC News API key goes here>")
