code: def get_stylize_fn():
      """Creates a tf.function for stylization."""
      input_spec = [
          tf.TensorSpec((None, None, None, 3), tf.float32),
          tf.TensorSpec((None, None, None, 3), tf.float32),
      ]
      predict_feeds = []
      predict_fetches = []
  
      def umbrella_function(content_img, style_img):
          predict_feeds.extend([content_img, style_img])
          predict_result = build_network(content_img, style_img)
          predict_fetches.extend(
              [
                  predict_result,
              ]
          )
          return predict_result
  
      umbrella_wrapped = tf.compat.v1.wrap_function(
          umbrella_function, input_spec
      )
      fn = umbrella_wrapped.prune(predict_feeds, predict_fetches)
      return fn
  
file_path: magenta/magenta/models/arbitrary_image_stylization/export_hub.py
imports: from absl import flags
  from magenta.models.arbitrary_image_stylization import arbitrary_image_stylization_build_model
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: get_stylize_fn() and returns: fn, predict_result