Update app.py

app.py

@@ -2,12 +2,11 @@ import tensorflow as tf
 import huggingface_hub as hf_hub
 import gradio as gr
 
-
 num_rows = 3
 num_cols = 3
 num_images = num_rows * num_cols
 image_size = 64
-plot_image_size = 64
+plot_image_size = 256
 
 model = hf_hub.from_pretrained_keras("beresandras/denoising-diffusion-model")
 
@@ -26,6 +25,7 @@ def generate_images(diffusion_steps, stochasticity, min_signal_rate, max_signal_
     step_size = 1.0 / diffusion_steps
     initial_noise = tf.random.normal(shape=(num_images, image_size, image_size, 3))
     
+    # reverse diffusion
     noisy_images = initial_noise
     for step in range(diffusion_steps):
         diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size
@@ -45,8 +45,14 @@ def generate_images(diffusion_steps, stochasticity, min_signal_rate, max_signal_
             + sample_noise_rates * sample_noises
         )
         
+    # denormalize
     generated_images = tf.clip_by_value(0.5 + 0.3 * pred_images, 0.0, 1.0)
+    
+    # make grid
     generated_images = tf.image.resize(generated_images, (plot_image_size, plot_image_size), method="nearest")
+    generated_images = tf.reshape(generated_images, (num_rows, num_cols, plot_image_size, plot_image_size, 3))
+    generated_images = tf.transpose(generated_images, (0, 2, 1, 3, 4))
+    generated_images = tf.reshape(generated_images, (num_rows * plot_image_size, num_cols * plot_image_size, 3))
     return generated_images.numpy()
 
 gr.Interface(