compile sampling step with autograph

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import tensorflow as tf
+import numpy as np
 from huggingface_hub import from_pretrained_keras
 from diffusion_sampler import DiffusionSampler
 
@@ -71,7 +72,7 @@ step_button = gr.Slider(
 )
 
 gallery = gr.Gallery(
-    columns=5,
+    columns=4,
     label="""
         Generated Flowers
     """
@@ -101,14 +102,16 @@ def call_model(
         progress=gr.Progress(track_tqdm=True),
 ):
     diffusion_model = linear_diffusion_model if model_to_call.lower() == "linear" else cosine_diffusion_model
-    return diffusion_model.generate_images(
-        num_images=int(num_images),
-        steps=int(steps),
+    images = diffusion_model.generate_images(
+        num_images=num_images,
+        steps=steps,
         sample_strategy=sample_strategy.lower(),
         step_strategy=step_strategy.lower(),
         ema=ema,
     )
 
+    return images.numpy().astype(np.uint8)
+
 
 demo = gr.Interface(
     fn=call_model,

diffusion_sampler.py

@@ -120,6 +120,7 @@ class DiffusionSampler(keras.Model):
         sqrt_one_minus_alpha_cum_prod = at_timestep(self.sqrt_one_minus_alphas_cum_prod)
         return sqrt_alpha_cum_prod * x_start + sqrt_one_minus_alpha_cum_prod * noise
 
+    @tf.function()
     def generate_images(
         self,
         num_images: int,
@@ -139,13 +140,9 @@ class DiffusionSampler(keras.Model):
         sampler, seq = sampling_stategies[(sample_strategy, step_strategy)]
         samples = tf.random.normal(shape=(num_images, 64, 64, 3), dtype=tf.float32)
 
-        for t in tqdm.tqdm(tf.reverse(seq, axis=[0])):
+        for t in tf.reverse(seq, axis=[0]):
             tt = tf.cast(tf.fill(dims=(num_images,), value=t), dtype=tf.int64)
-            pred_noise = noise_predictor.predict([samples, tt], verbose=0, batch_size=num_images)
+            pred_noise = noise_predictor([samples, tt], training=False)
             samples = sampler(pred_noise, samples, tt, )
 
-        return (
-            tf.clip_by_value(samples * 127.5 + 127.5, 0.0, 255.0)
-            .numpy()
-            .astype(np.uint8)
-        )
+        return tf.clip_by_value(samples * 127.5 + 127.5, 0.0, 255.0)