From b2cc256dabca983f417204ae3a2b382e6d1d5cc9 Mon Sep 17 00:00:00 2001
From: erogol <erogol@hotmail.com>
Date: Tue, 14 Jul 2020 17:48:44 +0200
Subject: [PATCH] tflite inference for melgan models

---
 vocoder/tf/models/melgan_generator.py         | 19 +++++++++++++++-
 .../tf/models/multiband_melgan_generator.py   | 22 ++++++++++++++-----
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/vocoder/tf/models/melgan_generator.py b/vocoder/tf/models/melgan_generator.py
index bf67f3d2..168fd29e 100644
--- a/vocoder/tf/models/melgan_generator.py
+++ b/vocoder/tf/models/melgan_generator.py
@@ -108,4 +108,21 @@ class MelganGenerator(tf.keras.models.Model):
 
     def build_inference(self):
         x = tf.random.uniform((1, self.in_channels, 4), dtype=tf.float32)
-        self(x, training=False)
\ No newline at end of file
+        self(x, training=False)
+
+    @tf.function(
+        experimental_relax_shapes=True,
+        input_signature=[
+            tf.TensorSpec([1, None, None], dtype=tf.float32),
+        ],)
+    def inference_tflite(self, c):
+        c = tf.transpose(c, perm=[0, 2, 1])
+        c = tf.expand_dims(c, 2)
+        # FIXME: TF had no replicate padding as in Torch
+        # c = tf.pad(c, [[0, 0], [self.inference_padding, self.inference_padding], [0, 0], [0, 0]], "REFLECT")
+        o = c
+        for layer in self.model_layers:
+            o = layer(o)
+        # o = self.model_layers(c)
+        o = tf.transpose(o, perm=[0, 3, 2, 1])
+        return o[:, :, 0, :]
\ No newline at end of file
diff --git a/vocoder/tf/models/multiband_melgan_generator.py b/vocoder/tf/models/multiband_melgan_generator.py
index c63ed06a..bdd333ed 100644
--- a/vocoder/tf/models/multiband_melgan_generator.py
+++ b/vocoder/tf/models/multiband_melgan_generator.py
@@ -30,11 +30,6 @@ class MultibandMelganGenerator(MelganGenerator):
     def pqmf_synthesis(self, x):
         return self.pqmf_layer.synthesis(x)
 
-    # def call(self, c, training=False):
-    #     if training:
-    #         raise NotImplementedError()
-    #     return self.inference(c)
-
     def inference(self, c):
         c = tf.transpose(c, perm=[0, 2, 1])
         c = tf.expand_dims(c, 2)
@@ -46,3 +41,20 @@ class MultibandMelganGenerator(MelganGenerator):
         o = tf.transpose(o, perm=[0, 3, 2, 1])
         o = self.pqmf_layer.synthesis(o[:, :, 0, :])
         return o
+
+    @tf.function(
+        experimental_relax_shapes=True,
+        input_signature=[
+            tf.TensorSpec([1, 80, None], dtype=tf.float32),
+        ],)
+    def inference_tflite(self, c):
+        c = tf.transpose(c, perm=[0, 2, 1])
+        c = tf.expand_dims(c, 2)
+        # FIXME: TF had no replicate padding as in Torch
+        # c = tf.pad(c, [[0, 0], [self.inference_padding, self.inference_padding], [0, 0], [0, 0]], "REFLECT")
+        o = c
+        for layer in self.model_layers:
+            o = layer(o)
+        o = tf.transpose(o, perm=[0, 3, 2, 1])
+        o = self.pqmf_layer.synthesis(o[:, :, 0, :])
+        return o