Add app.py

app.py

@@ -3,23 +3,32 @@ import numpy as np
 import jax
 import jax.numpy as jnp
 
-# Function that uses a seed to generate an image
+# Function to generate an image based on a seed
 def generate_image(seed):
-    key = jax.random.PRNGKey(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)
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Number(label="Seed", value=0),
-    outputs=gr.Image(type="numpy", label="Generated Image"),
-    title="JAX Diffusion Demo",
-    description="Enter a seed and click 'Generate' to sample a random diffusion image!",
-    allow_flagging="never",
-    examples=[[42], [1234], [2025]]  # Optional examples
-)
+# 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.")
+
+    with gr.Row():
+        seed_input = gr.Number(label="Seed", value=0)
+        generate_button = gr.Button("Generate")
+
+    output_image = gr.Image(label="Generated Image")
+
+    generate_button.click(
+        fn=generate_image,
+        inputs=seed_input,
+        outputs=output_image
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()
+