question-mark
Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

StableDiffusion Tensorflow to TF Lite

See original GitHub issue

Hi @LukeWood,

For fun, I tried converting stable Diffusion model from Tensorflow to TF lite, so that I can run it on coral/edge tpu.

I tried two approaches: I- Saved model approach: II- Go through h5

will try to document them as much as possible. (sorry in advance for the long traces)

for both:

!pip install git+https://github.com/divamgupta/stable-diffusion-tensorflow --upgrade 
!pip install tensorflow tensorflow_addons ftfy --upgrade

Using !pip install --upgrade keras-cv I was not able to save the model for both.

I- Saved model approach:

  1. Saved the model in a directory
from stable_diffusion_tf.stable_diffusion import StableDiffusion
model = StableDiffusion(
    img_height=512,
    img_width=512,
)
model.diffusion_model.save('/saved_model')
  1. lets try to load it:
import tensorflow as tf
model2 = tf.keras.models.load_model('/saved_model')
converter = tf.lite.TFLiteConverter.from_keras_model(model2)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)
  1. The following error is thrown:
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:absl:Found untraced functions such as dense_328_layer_call_fn, dense_328_layer_call_and_return_conditional_losses, dense_329_layer_call_fn, dense_329_layer_call_and_return_conditional_losses, group_normalization_173_layer_call_fn while saving (showing 5 of 1200). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: /tmp/tmpicdc9dmk/assets
INFO:tensorflow:Assets written to: /tmp/tmpicdc9dmk/assets
---------------------------------------------------------------------------
ConverterError                            Traceback (most recent call last)
<ipython-input-15-52d23a3e5390> in <module>
      2 model2 = tf.keras.models.load_model('mydata/ivo/pythalpha/saved_model')
      3 converter = tf.lite.TFLiteConverter.from_keras_model(model2)
----> 4 tflite_model = converter.convert()
      5 open("converted_model.tflite", "wb").write(tflite_model)

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/lite.py in wrapper(self, *args, **kwargs)
    931   def wrapper(self, *args, **kwargs):
    932     # pylint: disable=protected-access
--> 933     return self._convert_and_export_metrics(convert_func, *args, **kwargs)
    934     # pylint: enable=protected-access
    935 

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/lite.py in _convert_and_export_metrics(self, convert_func, *args, **kwargs)
    909     self._save_conversion_params_metric()
    910     start_time = time.process_time()
--> 911     result = convert_func(self, *args, **kwargs)
    912     elapsed_time_ms = (time.process_time() - start_time) * 1000
    913     if result:

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/lite.py in convert(self)
   1340         Invalid quantization parameters.
   1341     """
-> 1342     saved_model_convert_result = self._convert_as_saved_model()
   1343     if saved_model_convert_result:
   1344       return saved_model_convert_result

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/lite.py in _convert_as_saved_model(self)
   1322           self._convert_keras_to_saved_model(temp_dir))
   1323       if self.saved_model_dir:
-> 1324         return super(TFLiteKerasModelConverterV2,
   1325                      self).convert(graph_def, input_tensors, output_tensors)
   1326     finally:

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/lite.py in convert(self, graph_def, input_tensors, output_tensors)
   1133 
   1134     # Converts model.
-> 1135     result = _convert_graphdef(
   1136         input_data=graph_def,
   1137         input_tensors=input_tensors,

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/convert_phase.py in wrapper(*args, **kwargs)
    210         else:
    211           report_error_message(str(converter_error))
--> 212         raise converter_error from None  # Re-throws the exception.
    213       except Exception as error:
    214         report_error_message(str(error))

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/convert_phase.py in wrapper(*args, **kwargs)
    203     def wrapper(*args, **kwargs):
    204       try:
--> 205         return func(*args, **kwargs)
    206       except ConverterError as converter_error:
    207         if converter_error.errors:

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/convert.py in convert_graphdef(input_data, input_tensors, output_tensors, **kwargs)
    791       model_flags.output_arrays.append(util.get_tensor_name(output_tensor))
    792 
--> 793   data = convert(
    794       model_flags.SerializeToString(),
    795       conversion_flags.SerializeToString(),

/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/lite/python/convert.py in convert(model_flags_str, conversion_flags_str, input_data_str, debug_info_str, enable_mlir_converter)
    308       for error_data in _metrics_wrapper.retrieve_collected_errors():
    309         converter_error.append_error(error_data)
--> 310       raise converter_error
    311 
    312   return _run_deprecated_conversion_binary(model_flags_str,

ConverterError: /opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: error: 'tf.Conv2D' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["StatefulPartitionedCall:", "StatefulPartitionedCall"]): called from
/opt/conda/envs/rapids/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py:1267:0: note: Error code: ERROR_NEEDS_FLEX_OPS
<unknown>:0: error: failed while converting: 'main': 
Some ops are not supported by the native TFLite runtime, you can enable TF kernels fallback using TF Select. See instructions: https://www.tensorflow.org/lite/guide/ops_select 
TF Select ops: Conv2D
Details:
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<1x1x1280x1280xf32>) -> (tensor<?x?x?x1280xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<1x1x320x320xf32>) -> (tensor<?x?x?x320xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<1x1x640x640xf32>) -> (tensor<?x?x?x640xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x1280x1280xf32>) -> (tensor<?x?x?x1280xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x1280x640xf32>) -> (tensor<?x?x?x640xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x1920x1280xf32>) -> (tensor<?x?x?x1280xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x1920x640xf32>) -> (tensor<?x?x?x640xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x2560x1280xf32>) -> (tensor<?x?x?x1280xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x320x320xf32>) -> (tensor<?x?x?x320xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x320x4xf32>) -> (tensor<?x?x?x4xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x320x640xf32>) -> (tensor<?x?x?x640xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x640x1280xf32>) -> (tensor<?x?x?x1280xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x640x320xf32>) -> (tensor<?x?x?x320xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x640x640xf32>) -> (tensor<?x?x?x640xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x960x320xf32>) -> (tensor<?x?x?x320xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}
	tf.Conv2D(tensor<?x?x?x?xf32>, tensor<3x3x960x640xf32>) -> (tensor<?x?x?x640xf32>) : {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}

II- Go through h5

  1. Save the model with format h5
from stable_diffusion_tf.stable_diffusion import StableDiffusion
model = StableDiffusion(
    img_height=512,
    img_width=512,
)
model.diffusion_model.save('./stable_diffusion.h5', save_format='h5')
  1. Lets try to load it
import tensorflow as tf

model2 = tf.keras.models.load_model('stable_diffusion.h5')
converter = tf.lite.TFLiteConverter.from_keras_model(model2)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

  1. It seems that the TF 2.11.0 does not load h5 files anymore.
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[<ipython-input-3-72d9a214a713>](https://localhost:8080/#) in <module>
      1 import tensorflow as tf
      2 
----> 3 model2 = tf.keras.models.load_model('stable_diffusion.h5')
      4 converter = tf.lite.TFLiteConverter.from_keras_model(model2)
      5 tflite_model = converter.convert()

1 frames
[/usr/local/lib/python3.7/dist-packages/keras/saving/legacy/serialization.py](https://localhost:8080/#) in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    384     if cls is None:
    385         raise ValueError(
--> 386             f"Unknown {printable_module_name}: '{class_name}'. "
    387             "Please ensure you are using a `keras.utils.custom_object_scope` "
    388             "and that this object is included in the scope. See "

ValueError: Unknown layer: 'UNetModel'. Please ensure you are using a `keras.utils.custom_object_scope` and that this object is included in the scope. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.
  1. Therefore, uninstall tf 2.11.0 and install tf 2.1.0
  2. Attempt to load the saved h5 file:
import tensorflow as tf

model2 = tf.keras.models.load_model('stable_diffusion.h5')
converter = tf.lite.TFLiteConverter.from_keras_model(model2)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

The load_model throws the following error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-5-72d9a214a713> in <module>
      1 import tensorflow as tf
      2 
----> 3 model2 = tf.keras.models.load_model('stable_diffusion.h5')
      4 converter = tf.lite.TFLiteConverter.from_keras_model(model2)
      5 tflite_model = converter.convert()

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/saving/hdf5_format.py in load_model_from_hdf5(filepath, custom_objects, compile)
    164     if model_config is None:
    165       raise ValueError('No model found in config file.')
--> 166     model_config = json.loads(model_config.decode('utf-8'))
    167     model = model_config_lib.model_from_config(model_config,
    168                                                custom_objects=custom_objects)

AttributeError: 'str' object has no attribute 'decode'

Issue Analytics

  • State:open
  • Created 10 months ago
  • Reactions:2
  • Comments:38 (18 by maintainers)

github_iconTop GitHub Comments

4reactions
freedomtancommented, Dec 20, 2022

@charbull FYR. I gave a try.

ideally, if we feed input tensors with correct ranges, we should be able to get quantized model. if that works, we can extend the representative_data_gen().

  • input 1: 77x768 encoded text vectors, we can use the text encoder to generate encoded text prompt
  • input 2: I borrowed timestep embeddeding from the keras model
  • input 3: noise, since this is noise, I guess we can use random number
model = keras_cv.models.StableDiffusion(img_width=512, img_height=512)

prompt_1 = "A watercolor painting of a Golden Retriever at the beach"
encoding_1 = model.encode_text(prompt_1)


def get_timestep_embedding(timestep, batch_size, dim=320, max_period=10000):
    half = dim // 2
    freqs = tf.math.exp(
        -math.log(max_period) * tf.range(0, half, dtype=tf.float32) / half
    )

    args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
    embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
    embedding = tf.reshape(embedding, [1, -1])
    return tf.repeat(embedding, batch_size, axis=0)

def representative_data_gen():
  for i in range(1):
    em = get_timestep_embedding(i, 1)
    noise = tf.random.normal(1, 64, 64, 4)
    yield [encoding_1, em, noise]


converter = tf.lite.TFLiteConverter.from_keras_model(model.diffusion_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8  # or tf.uint8
converter.inference_output_type = tf.uint8  # or tf.uint8

converter.representative_dataset = representative_data_gen
tflite_model_int8 = converter.convert()

unfortunately, I got the following message. Dunno how to solve it yet.

Traceback (most recent call last):
  File "/hack/freedom/ml/notebooks/bar.py", line 45, in <module>
    tflite_model_int8 = converter.convert()
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 933, in wrapper
    return self._convert_and_export_metrics(convert_func, *args, **kwargs)
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 911, in _convert_and_export_metrics
    result = convert_func(self, *args, **kwargs)
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 1342, in convert
    saved_model_convert_result = self._convert_as_saved_model()
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 1325, in _convert_as_saved_model
    self).convert(graph_def, input_tensors, output_tensors)
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 1141, in convert
    return self._optimize_tflite_model(
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/convert_phase.py", line 215, in wrapper
    raise error from None  # Re-throws the exception.
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/convert_phase.py", line 205, in wrapper
    return func(*args, **kwargs)
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 871, in _optimize_tflite_model
    model = self._quantize(model, q_in_type, q_out_type, q_activations_type,
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/lite.py", line 608, in _quantize
    result = _calibrator.add_intermediate_tensors(result)
  File "/home/freedom/tf-master/lib/python3.10/site-packages/tensorflow/lite/python/optimize/calibrator.py", line 36, in add_intermediate_tensors
    return _calibration_wrapper.AddIntermediateTensors(model_content)
ValueError: Invalid model

3reactions
LukeWoodcommented, Nov 21, 2022

Thanks @charbull for the report! Will take a look at this. TFLite conversion would be an awesome addition.

Read more comments on GitHub >

github_iconTop Results From Across the Web

Stable Diffusion Tensorflow to TF Lite : r/StableDiffusion - Reddit
Hi, Checking here is someone tried to convert the tensorflow diffusion model into a tf lite?…
Read more >
High-performance image generation using Stable Diffusion in ...
Generate new images using KerasCV's StableDiffusion model. View on TensorFlow.org · Run in Google Colab · View source on GitHub · View on...
Read more >
divamgupta/stable-diffusion-tensorflow - TF Lite convert error
It appears that the problem is related to KerasCV's implementation because they are using their own implementation of the GroupNormalization layer, which is ......
Read more >
Anyone attempted to convert stablediffusion tensorflow to tf lite?
was curious if someone attempted the conversion? I tried here github.com/divamgupta/stable-diffusion-tensorflow/issues/58 but having some ...
Read more >
How to Run TensorFlow Lite Models on Raspberry Pi
In this tutorial we'll see how to run TensorFlow Lite on Raspberry Pi. We'll use the TFLite version of MobileNet for making predictions...
Read more >

github_iconTop Related Medium Post

No results found

github_iconTop Related StackOverflow Question

No results found

github_iconTroubleshoot Live Code

Lightrun enables developers to add logs, metrics and snapshots to live code - no restarts or redeploys required.
Start Free

github_iconTop Related Reddit Thread

No results found

github_iconTop Related Hackernoon Post

No results found

github_iconTop Related Tweet

No results found

github_iconTop Related Dev.to Post

No results found

github_iconTop Related Hashnode Post

No results found