Updated Layout

app.py

@@ -4,15 +4,11 @@ from PIL import Image
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
 from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
+from diffusers.utils import load_image
 import jax.numpy as jnp
 import numpy as np
 
 
-title = "🧨 ControlNet on Segment Anything 🤗"
-description = "This is a demo on ControlNet based on Segment Anything"
-
-examples = [["a modern main room of a house", "low quality", "condition_image_1.png", 50, 4]]
-
 controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
     "mfidabel/controlnet-segment-anything", dtype=jnp.float32
 )
@@ -25,13 +21,18 @@ pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
 params["controlnet"] = controlnet_params
 p_params = replicate(params)
 
+# Description
+title = "# 🧨 ControlNet on Segment Anything 🤗"
+description = "This is a demo on ControlNet based on Segment Anything"
+
+examples = [["a modern main room of a house", "low quality", "condition_image_1.png", 50, 4, 4]]
 
 # Inference Function
-def infer(prompts, negative_prompts, image, num_inference_steps, seed):
+def infer(prompts, negative_prompts, image, num_inference_steps, seed, num_samples):
     rng = jax.random.PRNGKey(int(seed))
     num_inference_steps = int(num_inference_steps)
     image = Image.fromarray(image, mode="RGB")
-    num_samples = jax.device_count()
+    num_samples = max(jax.device_count(), int(num_samples))
     p_rng = jax.random.split(rng, jax.device_count())
     
     prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)
@@ -59,10 +60,45 @@ def infer(prompts, negative_prompts, image, num_inference_steps, seed):
     del output
     
     return final_image
+
+with gr.Blocks(css="h1 { text-align: center }") as demo:
+    # Title
+    gr.Markdown(title)
+    # Description
+    gr.Markdown(description)
+
+    # Images
+    with gr.Row(variant="panel"):
+        cond_img = gr.Image(label="Input")\
+                     .style(height=400)
+        output = gr.Gallery(label="Generated images")\
+                   .style(height=400, rows=[2], columns=[2])
+        
+    # Submit & Clear
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(lines=1, label="Prompt")
+            negative_prompt = gr.Textbox(lines=1, label="Negative Prompt")
+
+        with gr.Column():
+            with gr.Accordion("Advanced options", open=False):
+                num_steps = gr.Slider(10, 60, 50, step=1, label="Steps")
+                seed = gr.Slider(0, 1024, 0, step=1, label="Seed")
+                num_samples = gr.Slider(1, 4, 4, step=1, label="Nº Samples")
+                
+            submit = gr.Button("Submit")
+
+    # Examples
+    gr.Examples(examples=examples,
+                inputs=[prompt, negative_prompt, cond_img, num_steps, seed, num_samples],
+                outputs=output,
+                fn=infer,
+                cache_examples=True)
+
+    
+    submit.click(infer, 
+                 inputs=[prompt, negative_prompt, cond_img, num_steps, seed, num_samples],
+                 outputs = output)
     
-gr.Interface(fn = infer, 
-             inputs = ["text", "text", "image", "number", "number"],
-             outputs = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto", preview=True),
-             title = title,
-             description = description,
-             examples = examples).launch()
+demo.queue()
+demo.launch()