code: def _export_to_saved_model(checkpoint, alpha, num_styles):
      """Export a image stylization checkpoint to SavedModel format."""
      saved_model_dir = tempfile.mkdtemp()
  
      with tf.Graph().as_default(), tf.Session() as sess:
          # Define input tensor as placeholder to allow export
          input_image_tensor = tf.placeholder(
              tf.float32, shape=(1, None, None, 3), name="input_image"
          )
          weights_tensor = tf.placeholder(
              tf.float32, shape=num_styles, name="style_weights"
          )
  
          # Load the graph definition from Magenta
          stylized_image_tensor = model.transform(
              input_image_tensor,
              alpha=alpha,
              normalizer_fn=ops.weighted_instance_norm,
              normalizer_params={
                  "weights": weights_tensor,
                  "num_categories": FLAGS.num_styles,
                  "center": True,
                  "scale": True,
              },
          )
  
          # Load model weights from downloaded checkpoint file
          _load_checkpoint(sess, checkpoint)
  
          # Write SavedModel for serving or conversion to TF Lite
          tf.saved_model.simple_save(
              sess,
              saved_model_dir,
              inputs={
                  input_image_tensor.name: input_image_tensor,
                  weights_tensor.name: weights_tensor,
              },
              outputs={"stylized_image": stylized_image_tensor},
          )
  
      return saved_model_dir
  
file_path: magenta/magenta/models/image_stylization/image_stylization_convert_tflite.py
imports: import os
  import tempfile
  from magenta.models.image_stylization import model
  from magenta.models.image_stylization import ops
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: _export_to_saved_model(checkpoint, alpha, num_styles) and returns: saved_model_dir