refactor code

README.md

@@ -1,13 +1,15 @@
 ---
-title: Temp Diffusion
-emoji: 📈
-colorFrom: indigo
-colorTo: pink
+title: Diffusion Model
+emoji: 🌼
+colorFrom: yellow
+colorTo: orange
 sdk: gradio
 sdk_version: 4.10.0
 app_file: app.py
-pinned: false
+pinned: true
 license: agpl-3.0
+suggested_storage: "small"
+suggested_hardware: "t4-small"
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,11 +1,8 @@
 import gradio as gr
-import tensorflow as tf
 import numpy as np
 from huggingface_hub import from_pretrained_keras
 from diffusion_sampler import DiffusionSampler
 
-print(f"detected GPUs={tf.config.list_physical_devices('GPU')}")
-
 scheduler_button = gr.Radio(
     choices=["Linear", "Cosine"],
     label="Noise Scheduler",
@@ -81,18 +78,21 @@ gallery = gr.Gallery(
     """
 )
 
+diffusion_models = {
+    "linear":
+        DiffusionSampler(
+            model=from_pretrained_keras("leowajda/linear_diffusion"),
+            ema_model=from_pretrained_keras("leowajda/linear_diffusion_ema"),
+            noise_scheduler="linear"
+        ),
 
-linear_diffusion_model = DiffusionSampler(
-    model=from_pretrained_keras("leowajda/linear_diffusion"),
-    ema_model=from_pretrained_keras("leowajda/linear_diffusion_ema"),
-    noise_scheduler="linear",
-)
-
-cosine_diffusion_model = DiffusionSampler(
-    model=from_pretrained_keras("leowajda/cosine_diffusion"),
-    ema_model=from_pretrained_keras("leowajda/cosine_diffusion_ema"),
-    noise_scheduler="cosine",
-)
+    "cosine":
+        DiffusionSampler(
+            model=from_pretrained_keras("leowajda/cosine_diffusion"),
+            ema_model=from_pretrained_keras("leowajda/cosine_diffusion_ema"),
+            noise_scheduler="cosine"
+        )
+}
 
 
 def call_model(
@@ -102,9 +102,8 @@ def call_model(
         ema: bool = True,
         steps: int = 1_000,
         num_images: int = 0,
-        progress=gr.Progress(track_tqdm=True),
 ):
-    diffusion_model = linear_diffusion_model if model_to_call.lower() == "linear" else cosine_diffusion_model
+    diffusion_model = diffusion_models[model_to_call.lower()]
     images = diffusion_model.generate_images(
         num_images=num_images,
         steps=steps,

diffusion_sampler.py

@@ -1,5 +1,3 @@
-import numpy as np
-import tqdm as tqdm
 import tensorflow as tf
 import math
 from tensorflow import keras
@@ -100,10 +98,10 @@ class DiffusionSampler(keras.Model):
         return sqrt_alpha_cum_prod_prev * x0_t + x_t_dir * pred_noise + c1 * random_noise
 
     def noise_scheduler(self, scheduler: str, max_beta: int = 0.02) -> tf.Tensor:
-        alpha_bar = lambda t: tf.math.cos((t + 0.008) / 1.008 * tf.constant(math.pi, dtype=tf.float64) / 2) ** 2
-        cosine_scheduler = lambda i: tf.minimum(
-            1 - alpha_bar((i + 1) / tf.cast(self.timesteps, dtype=tf.float64)) / alpha_bar(
-                i / tf.cast(self.timesteps, dtype=tf.float64)), max_beta)
+        pi, t = [tf.constant(num, dtype=tf.float64) for num in (math.pi, self.timesteps)]
+
+        alpha_bar = lambda t: tf.math.cos((t + 0.008) / 1.008 * pi / 2) ** 2
+        cosine_scheduler = lambda i: tf.minimum(1 - alpha_bar((i + 1) / t) / alpha_bar(i / t), max_beta)
 
         if scheduler == "linear":
             x = tf.linspace(start=self.beta_start, stop=self.beta_end, num=self.timesteps)