Update app.py

app.py

@@ -1,29 +1,33 @@
-# app.py
-
 import gradio as gr
 import jax
 import jax.numpy as jnp
-import numpy as np
-from PIL import Image
 
-# Dummy generator function — Replace this with your real model inference!
-def generate_image(seed):
-    key = jax.random.PRNGKey(seed)
-    # Generate a fake "image" of size 64x64x3 (RGB)
-    img = jax.random.uniform(key, (64, 64, 3), minval=0, maxval=1.0)
-    img_np = np.array(img * 255, dtype=np.uint8)
-    return Image.fromarray(img_np)
+# Dummy diffusion generation function (replace with your real one)
+def generate_diffusion(prompt, steps):
+    key = jax.random.PRNGKey(0)
+    # For demo: Create random noise image
+    image = jax.random.uniform(key, (64, 64, 3))
+    image = jnp.clip(image, 0, 1)
+    return image
+
+# Gradio Interface using Blocks
+with gr.Blocks() as demo:
+    gr.Markdown("# 🌟 JAX Diffusion Demo")
+    gr.Markdown("Generate images using a simple diffusion model powered by **JAX**!")
+
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Describe your image...")
+            steps_input = gr.Slider(10, 100, value=50, step=5, label="Diffusion Steps")
+            generate_button = gr.Button("Generate")
+
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
 
-# Define Gradio Interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Slider(0, 10000, value=42, step=1, label="Random Seed"),
-    outputs=gr.Image(type="pil", label="Generated Image"),
-    title="JAX Diffusion Demo",
-    description="🎨 Generate random diffusion samples using JAX! \n\n(Replace dummy function with your trained model.)",
-    theme="default",
-    live=False
-)
+    generate_button.click(
+        fn=generate_diffusion,
+        inputs=[prompt_input, steps_input],
+        outputs=output_image
+    )
 
-if __name__ == "__main__":
-    iface.launch()
+demo.launch()