Create app.py

app.py

@@ -0,0 +1,70 @@
+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
+
+
+def load_model():
+    model = hf_hub.from_pretrained_keras("beresandras/denoising-diffusion-model")
+    return model
+
+def diffusion_schedule(diffusion_times, min_signal_rate, max_signal_rate):
+    start_angle = tf.acos(max_signal_rate)
+    end_angle = tf.acos(min_signal_rate)
+
+    diffusion_angles = start_angle + diffusion_times * (end_angle - start_angle)
+
+    signal_rates = tf.cos(diffusion_angles)
+    noise_rates = tf.sin(diffusion_angles)
+    
+    return noise_rates, signal_rates
+
+def generate_images(model, num_images, diffusion_steps, stochasticity, min_signal_rate, max_signal_rate):
+    step_size = 1.0 / diffusion_steps
+    initial_noise = tf.random.normal(shape=(num_images, image_size, image_size, 3))
+    
+    noisy_images = initial_noise
+    for step in range(diffusion_steps):
+        diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size
+        next_diffusion_times = diffusion_times - step_size
+        
+        noise_rates, signal_rates = diffusion_schedule(diffusion_times, min_signal_rate, max_signal_rate)
+        next_noise_rates, next_signal_rates = diffusion_schedule(next_diffusion_times, min_signal_rate, max_signal_rate)
+        
+        sample_noises = tf.random.normal(shape=(num_images, image_size, image_size, 3))
+        sample_noise_rates = stochasticity * (1.0 - (signal_rates / next_signal_rates)**2)**0.5 * (next_noise_rates / noise_rates)
+        
+        pred_noises = model([noisy_images, noise_rates])
+        pred_images = (noisy_images - noise_rates * pred_noises) / signal_rates
+        noisy_images = (
+            next_signal_rates * pred_images
+            + (next_noise_rates**2 - sample_noise_rates**2)**0.5 * pred_noises
+            + sample_noise_rates * sample_noises
+        )
+        
+    generated_images = tf.clip_by_value(0.5 + 0.3 * pred_images, 0.0, 1.0)
+    generated_images = tf.image.resize(
+            generated_images, (plot_image_size, plot_image_size), method="nearest"
+    )
+    return generated_images.numpy()
+
+
+model = load_model()
+gr.Interface(
+    generate_images,
+    inputs=[
+        model,
+        num_images,
+        gr.inputs.Slider(1, 20, default=10, label="Diffusion steps"),
+        gr.inputs.Slider(0.0, 1.0, step=0.05, default=0.0, label="Stochasticity"),
+        gr.inputs.Slider(0.02, 0.10, step=0.01, default=0.02, label="Minimal signal rate"),
+        gr.inputs.Slider(0.80, 0.95, step=0.01, default=0.95, label="Maximal signal rate"),
+    ],
+    outputs="image",
+).launch()