Add app.py

app.py

@@ -1,34 +1,30 @@
-import gradio as gr
-import numpy as np
 import jax
 import jax.numpy as jnp
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
 
-# Function to generate an image based on a seed
+# Dummy diffusion sampling function
 def generate_image(seed):
-    if seed is None:
-        seed = 0  # default fallback
-    key = jax.random.PRNGKey(int(seed))
-    img = jax.random.normal(key, (64, 64, 3))
-    img = (img - img.min()) / (img.max() - img.min())  # Normalize to [0, 1]
-    return np.array(img)
+    key = jax.random.PRNGKey(seed)
+    img = jax.random.uniform(key, (64, 64, 3))
+    img = np.array(img)  # Convert to numpy for display
+    return img
 
-# Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🌟 JAX Diffusion Demo")
-    gr.Markdown("Enter a random seed to generate a diffusion-based image.")
+    gr.Markdown("# 🌀 JAX Diffusion Demo")
+    gr.Markdown("Generate random images using JAX diffusion (dummy example).")
 
     with gr.Row():
-        seed_input = gr.Number(label="Seed", value=0)
-        generate_button = gr.Button("Generate")
-
-    output_image = gr.Image(label="Generated Image")
+        seed_slider = gr.Slider(minimum=0, maximum=10000, step=1, value=42, label="Seed")
+        generate_button = gr.Button("Generate Image")
+    
+    output_image = gr.Image(type="numpy", label="Generated Image", shape=(64, 64))
 
     generate_button.click(
         fn=generate_image,
-        inputs=seed_input,
+        inputs=seed_slider,
         outputs=output_image
     )
 
-if __name__ == "__main__":
-    demo.launch()
-
+demo.launch()