code: def predict_model_gen(session, style_dataset, sample_count):
      """Create a generator function that emits style images.
  
      Args:
        session: tf.Session, the session that contains subgraph to load the training
          dataset
        style_dataset: tf.data.Dataset that contains training style images.
        sample_count: int, number of sample to create.
  
      Returns:
        (str, str) A generator function to use as representative dataset for
        TFLiteConverter.
      """
  
      def generator():
          dataset = style_dataset.batch(1)
          iterator = tf.data.make_initializable_iterator(dataset)
          session.run(iterator.initializer)
          next_element = iterator.get_next()
          for _ in range(sample_count):
              input_value = session.run(next_element)
              yield [input_value]
  
      return generator
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import tempfile
  from magenta.models.arbitrary_image_stylization import arbitrary_image_stylization_build_mobilenet_model
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 8 for calling a global function with example usage: predict_model_gen(session, style_dataset, sample_count) and returns: generator