Merge pull request #153 from MahmoudAshraf97:master

PiperOrigin-RevId: 487567856
Change-Id: I82dc05d4229c561045f4160c4aefda334f212e91

Author: Johannes Ballé

models/hific/model.py

@@ -846,14 +846,15 @@ def __init__(self, weight_path):
     def wrap_frozen_graph(graph_def, inputs, outputs):
       def _imports_graph_def():
         tf.graph_util.import_graph_def(graph_def, name="")
-      wrapped_import = tf.wrap_function(_imports_graph_def, [])
+
+      wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
       import_graph = wrapped_import.graph
       return wrapped_import.prune(
           tf.nest.map_structure(import_graph.as_graph_element, inputs),
           tf.nest.map_structure(import_graph.as_graph_element, outputs))
 
     # Pack LPIPS network into a tf function
-    graph_def = tf.GraphDef()
+    graph_def = tf.compat.v1.GraphDef()
     with open(weight_path, "rb") as f:
       graph_def.ParseFromString(f.read())
     self._lpips_func = tf.function(