Add app.py

app.py

@@ -3,9 +3,9 @@ import numpy as np
 import jax
 import jax.numpy as jnp
 
-# Dummy simple function to simulate diffusion sampling
-def generate_image():
-    key = jax.random.PRNGKey(np.random.randint(0, 10000))
+# Function that uses a seed to generate an image
+def generate_image(seed):
+    key = jax.random.PRNGKey(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)
@@ -13,11 +13,12 @@ def generate_image():
 # Gradio Interface
 iface = gr.Interface(
     fn=generate_image,
-    inputs=[],
+    inputs=gr.Number(label="Seed", value=0),
     outputs=gr.Image(type="numpy", label="Generated Image"),
     title="JAX Diffusion Demo",
-    description="Click 'Generate' to sample a random diffusion image!",
-    allow_flagging="never"
+    description="Enter a seed and click 'Generate' to sample a random diffusion image!",
+    allow_flagging="never",
+    examples=[[42], [1234], [2025]]  # Optional examples
 )
 
 if __name__ == "__main__":