code: def convert_unet(pipe, args, model_name=None):
      """Converts the UNet component of Stable Diffusion"""
      if args.unet_support_controlnet:
          unet_name = "control-unet"
      else:
          unet_name = model_name or "unet"
  
      out_path = _get_out_path(args, unet_name)
  
      # Check if Unet was previously exported and then chunked
      unet_chunks_exist = all(
          os.path.exists(
              out_path.replace(".mlpackage", f"_chunk{idx+1}.mlpackage")
          )
          for idx in range(2)
      )
  
      if args.chunk_unet and unet_chunks_exist:
          logger.info(
              "`unet` chunks already exist, skipping conversion."
          )
          del pipe.unet
          gc.collect()
          return
  
      # If original Unet does not exist, export it from PyTorch+diffusers
      elif not os.path.exists(out_path):
          # Prepare sample input shapes and values
          batch_size = 2  # for classifier-free guidance
          sample_shape = (
              batch_size,  # B
              pipe.unet.config.in_channels,  # C
              args.latent_h or pipe.unet.config.sample_size,  # H
              args.latent_w or pipe.unet.config.sample_size,  # W
          )
  
          if not hasattr(pipe, "text_encoder"):
              raise RuntimeError(
                  "convert_text_encoder() deletes pipe.text_encoder to save RAM. "
                  "Please use convert_unet() before convert_text_encoder()"
              )
  
          if (
              hasattr(pipe, "text_encoder")
              and pipe.text_encoder is not None
          ):
              text_token_sequence_length = (
                  pipe.text_encoder.config.max_position_embeddings
              )
              hidden_size = (pipe.text_encoder.config.hidden_size,)
          elif (
              hasattr(pipe, "text_encoder_2")
              and pipe.text_encoder_2 is not None
          ):
              text_token_sequence_length = (
                  pipe.text_encoder_2.config.max_position_embeddings
              )
              hidden_size = (pipe.text_encoder_2.config.hidden_size,)
  
          encoder_hidden_states_shape = (
              batch_size,
              args.text_encoder_hidden_size
              or pipe.unet.config.cross_attention_dim
              or hidden_size,
              1,
              args.text_token_sequence_length
              or text_token_sequence_length,
          )
  
          # Create the scheduled timesteps for downstream use
          DEFAULT_NUM_INFERENCE_STEPS = 50
          pipe.scheduler.set_timesteps(DEFAULT_NUM_INFERENCE_STEPS)
  
          sample_unet_inputs = OrderedDict(
              [
                  ("sample", torch.rand(*sample_shape)),
                  (
                      "timestep",
                      torch.tensor(
                          [pipe.scheduler.timesteps[0].item()]
                          * (batch_size)
                      ).to(torch.float32),
                  ),
                  (
                      "encoder_hidden_states",
                      torch.rand(*encoder_hidden_states_shape),
                  ),
              ]
          )
  
          # Prepare inputs
          baseline_sample_unet_inputs = deepcopy(sample_unet_inputs)
          baseline_sample_unet_inputs["encoder_hidden_states"] = (
              baseline_sample_unet_inputs["encoder_hidden_states"]
              .squeeze(2)
              .transpose(1, 2)
          )
  
          # Initialize reference unet
          if args.xl_version:
              unet_cls = unet.UNet2DConditionModelXL
  
              # Sample time_ids
              height = (
                  args.latent_h or pipe.unet.config.sample_size
              ) * 8
              width = (
                  args.latent_w or pipe.unet.config.sample_size
              ) * 8
  
              original_size = (height, width)  # output_resolution
              crops_coords_top_left = (0, 0)  # topleft_crop_cond
              target_size = (height, width)  # resolution_cond
              if (
                  hasattr(pipe.config, "requires_aesthetics_score")
                  and pipe.config.requires_aesthetics_score
              ):
                  # Part of SDXL's micro-conditioning as explained in section 2.2 of
                  # [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to
                  # simulate an aesthetic score of the generated image by influencing the positive and negative text conditions.
                  aesthetic_score = 6.0  # default aesthetic_score
                  negative_aesthetic_score = (
                      2.5  # default negative_aesthetic_score
                  )
                  add_time_ids = list(
                      original_size
                      + crops_coords_top_left
                      + (aesthetic_score,)
                  )
                  add_neg_time_ids = list(
                      original_size
                      + crops_coords_top_left
                      + (negative_aesthetic_score,)
                  )
              else:
                  add_time_ids = list(
                      original_size
                      + crops_coords_top_left
                      + target_size
                  )
                  add_neg_time_ids = list(
                      original_size
                      + crops_coords_top_left
                      + target_size
                  )
  
              time_ids = [add_neg_time_ids, add_time_ids]
  
              # Pooled text embedding from text_encoder_2
              text_embeds_shape = (
                  batch_size,
                  pipe.text_encoder_2.config.hidden_size,
              )
  
              additional_xl_inputs = OrderedDict(
                  [
                      (
                          "time_ids",
                          torch.tensor(time_ids).to(torch.float32),
                      ),
                      ("text_embeds", torch.rand(*text_embeds_shape)),
                  ]
              )
  
              sample_unet_inputs.update(additional_xl_inputs)
              baseline_sample_unet_inputs[
                  "added_cond_kwargs"
              ] = additional_xl_inputs
          else:
              unet_cls = unet.UNet2DConditionModel
  
          reference_unet = unet_cls(**pipe.unet.config).eval()
  
          load_state_dict_summary = reference_unet.load_state_dict(
              pipe.unet.state_dict()
          )
  
          if args.unet_support_controlnet:
              from .unet import calculate_conv2d_output_shape
  
              additional_residuals_shapes = []
  
              # conv_in
              out_h, out_w = calculate_conv2d_output_shape(
                  (args.latent_h or pipe.unet.config.sample_size),
                  (args.latent_w or pipe.unet.config.sample_size),
                  reference_unet.conv_in,
              )
              additional_residuals_shapes.append(
                  (
                      batch_size,
                      reference_unet.conv_in.out_channels,
                      out_h,
                      out_w,
                  )
              )
  
              # down_blocks
              for down_block in reference_unet.down_blocks:
                  additional_residuals_shapes += [
                      (batch_size, resnet.out_channels, out_h, out_w)
                      for resnet in down_block.resnets
                  ]
                  if (
                      hasattr(down_block, "downsamplers")
                      and down_block.downsamplers is not None
                  ):
                      for downsampler in down_block.downsamplers:
                          out_h, out_w = calculate_conv2d_output_shape(
                              out_h, out_w, downsampler.conv
                          )
                      additional_residuals_shapes.append(
                          (
                              batch_size,
                              down_block.downsamplers[
                                  -1
                              ].conv.out_channels,
                              out_h,
                              out_w,
                          )
                      )
  
              # mid_block
              additional_residuals_shapes.append(
                  (
                      batch_size,
                      reference_unet.mid_block.resnets[-1].out_channels,
                      out_h,
                      out_w,
                  )
              )
  
              baseline_sample_unet_inputs[
                  "down_block_additional_residuals"
              ] = ()
              for i, shape in enumerate(additional_residuals_shapes):
                  sample_residual_input = torch.rand(*shape)
                  sample_unet_inputs[
                      f"additional_residual_{i}"
                  ] = sample_residual_input
                  if i == len(additional_residuals_shapes) - 1:
                      baseline_sample_unet_inputs[
                          "mid_block_additional_residual"
                      ] = sample_residual_input
                  else:
                      baseline_sample_unet_inputs[
                          "down_block_additional_residuals"
                      ] += (sample_residual_input,)
  
          sample_unet_inputs_spec = {
              k: (v.shape, v.dtype)
              for k, v in sample_unet_inputs.items()
          }
          logger.info(
              f"Sample UNet inputs spec: {sample_unet_inputs_spec}"
          )
  
          # JIT trace
          logger.info("JIT tracing..")
          reference_unet = torch.jit.trace(
              reference_unet, list(sample_unet_inputs.values())
          )
          logger.info("Done.")
  
          if args.check_output_correctness:
              baseline_out = pipe.unet.to(torch.float32)(
                  **baseline_sample_unet_inputs, return_dict=False
              )[0].numpy()
              reference_out = reference_unet(
                  *sample_unet_inputs.values()
              )[0].numpy()
              report_correctness(
                  baseline_out,
                  reference_out,
                  "unet baseline to reference PyTorch",
              )
  
          del pipe.unet
          gc.collect()
  
          coreml_sample_unet_inputs = {
              k: v.numpy().astype(np.float16)
              for k, v in sample_unet_inputs.items()
          }
  
          coreml_unet, out_path = _convert_to_coreml(
              unet_name,
              reference_unet,
              coreml_sample_unet_inputs,
              ["noise_pred"],
              args,
          )
          del reference_unet
          gc.collect()
  
          # Set model metadata
          coreml_unet.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
          if args.xl_version:
              coreml_unet.license = "OpenRAIL++-M (https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md)"
          else:
              coreml_unet.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
          coreml_unet.version = (
              args.model_version
              if model_name != "refiner"
              or not hasattr(args, "refiner_version")
              else args.refiner_version
          )
          coreml_unet.short_description = (
              "Stable Diffusion generates images conditioned on text or other images as input through the diffusion process. "
              "Please refer to https://arxiv.org/abs/2112.10752 for details."
          )
  
          # Set the input descriptions
          coreml_unet.input_description[
              "sample"
          ] = "The low resolution latent feature maps being denoised through reverse diffusion"
          coreml_unet.input_description[
              "timestep"
          ] = "A value emitted by the associated scheduler object to condition the model on a given noise schedule"
          coreml_unet.input_description["encoder_hidden_states"] = (
              "Output embeddings from the associated text_encoder model to condition to generated image on text. "
              "A maximum of 77 tokens (~40 words) are allowed. Longer text is truncated. "
              "Shorter text does not reduce computation."
          )
          if args.xl_version:
              coreml_unet.input_description[
                  "time_ids"
              ] = "Additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks."
              coreml_unet.input_description[
                  "text_embeds"
              ] = "Additional embeddings from text_encoder_2 that if specified are added to the embeddings that are passed along to the UNet blocks."
  
          # Set the output descriptions
          coreml_unet.output_description["noise_pred"] = (
              "Same shape and dtype as the `sample` input. "
              "The predicted noise to facilitate the reverse diffusion (denoising) process"
          )
  
          # Set package version metadata
          from python_coreml_stable_diffusion._version import (
              __version__,
          )
  
          coreml_unet.user_defined_metadata[
              "com.github.apple.ml-stable-diffusion.version"
          ] = __version__
  
          coreml_unet.save(out_path)
          logger.info(f"Saved unet into {out_path}")
  
          # Parity check PyTorch vs CoreML
          if args.check_output_correctness:
              coreml_out = list(
                  coreml_unet.predict(
                      coreml_sample_unet_inputs
                  ).values()
              )[0]
              report_correctness(
                  baseline_out,
                  coreml_out,
                  "unet baseline PyTorch to reference CoreML",
              )
  
          del coreml_unet
          gc.collect()
      else:
          del pipe.unet
          gc.collect()
          logger.info(
              f"`unet` already exists at {out_path}, skipping conversion."
          )
  
      if args.chunk_unet and not unet_chunks_exist:
          logger.info(
              f"Chunking {model_name} in two approximately equal MLModels"
          )
          args.mlpackage_path = out_path
          args.remove_original = False
          args.merge_chunks_in_pipeline_model = False
          chunk_mlprogram.main(args)
  
imports: from python_coreml_stable_diffusion import unet, controlnet, chunk_mlprogram
  import argparse
  from collections import OrderedDict, defaultdict
  from copy import deepcopy
  import coremltools
  from diffusers import StableDiffusionPipeline, DiffusionPipeline, ControlNetModel
  import gc
  import logging
  import numpy
  import os
  import requests
  import shutil
  import time
  import re
  import pathlib
  import torch
  import torch.nn
  import torch.nn.functional
  from types import MethodType
  from transformers.models.clip import modeling_clip
label_desc: Use image node_id 15 for calling a global function with example usage: convert_unet(pipe, args, model_name) without return types