minor improvements

.gitignore

@@ -1 +1 @@
-.idea
+.idea/
 __pycache__/

README.md

@@ -1,6 +1,6 @@
 ---
 title: Diffusion Model
-emoji: 🌼
+emoji: 🌼🌸🤗🏵️🌺
 colorFrom: yellow
 colorTo: red
 sdk: gradio

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
 import numpy as np
+import tensorflow as tf
 from huggingface_hub import from_pretrained_keras
 from diffusion_sampler import DiffusionSampler
 
+print('\n'.join([f'- {device.name}' for device in tf.config.list_physical_devices('GPU')]) or 'No GPU devices found.')
+
 scheduler_button = gr.Radio(
     choices=["Linear", "Cosine"],
     label="Noise Scheduler",
@@ -39,8 +42,7 @@ ema_button = gr.Checkbox(
     value=True,
     label="Exponential Moving Average",
     info="""
-        Whether to invoke the network with the applied exponential moving average on the model parameters.
-        Recommended for better results.
+        Whether to invoke the network with the applied exponential moving average on the model's weights.
     """
 )
 
@@ -57,7 +59,7 @@ images_button = gr.Number(
 )
 
 step_button = gr.Slider(
-    minimum=500,
+    minimum=700,
     value=1_000,
     maximum=1_000,
     randomize=True,
@@ -119,12 +121,12 @@ demo = gr.Interface(
     fn=call_model,
     inputs=[scheduler_button, sampling_button, subsequence_button, ema_button, step_button, images_button],
     outputs=gallery,
-    cache_examples=False,
+    cache_examples=True,
     title="""Unconditional Image Generation Through Denoising Diffusion Implicit Models""",
     examples=[
         ["Linear", "DDPM", "Linear", True, 1_000, 10],
-        ["Cosine", "DDIM", "Linear", True, 750, 20],
-        ["Linear", "DDIM", "Quadratic", True, 750, 20]
+        ["Cosine", "DDPM", "Linear", True, 750, 20],
+        ["Linear", "DDIM", "Linear", True, 750, 20]
     ],
     description="""                
         <p align="center">

diffusion_sampler.py

@@ -16,7 +16,7 @@ def batch_reshape(t: tf.Tensor, x: tf.Tensor) -> tf.Tensor:
     return inner_function
 
 
-class DiffusionSampler(keras.Model):
+class DiffusionSampler:
     def __init__(
         self,
         model: keras.Model | str,
@@ -26,9 +26,7 @@ class DiffusionSampler(keras.Model):
         beta_end: float | None = 0.02,
         noise_scheduler: str = "linear",
         ema: float = 0.999,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
         self.noise_predictor = load_model(filepath=model, safe_mode=False) if isinstance(model, str) else model
         self.ema_noise_predictor = load_model(filepath=ema_model, safe_mode=False) if isinstance(model, str) else ema_model
         self.ema = ema
@@ -118,7 +116,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(reduce_retracing=True)
+    @tf.function
     def generate_images(
         self,
         num_images: int,