Changed Example Layout, Predefined Input

app.py

@@ -7,6 +7,7 @@ from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
 from diffusers.utils import load_image
 import jax.numpy as jnp
 import numpy as np
+import gc
 
 
 controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
@@ -31,67 +32,94 @@ description = """This is a demo on 🧨 ControlNet based on Meta's [Segment Anyt
                 Test some of the examples below to give it a try ⬇️
               """
 
-examples = [["a modern main room of a house", "low quality", "examples/condition_image_1.png"],
+examples = [["contemporary living room of a house", "low quality", "examples/condition_image_1.png"],
             ["new york buildings,  Vincent Van Gogh starry night ", "low quality, monochrome", "examples/condition_image_2.png"],
             ["contemporary living room,  high quality, 4k, realistic", "low quality, monochrome, low res", "examples/condition_image_3.png"]]
 
 
 # Inference Function
 def infer(prompts, negative_prompts, image, num_inference_steps = 50, seed = 4, num_samples = 4):
-    rng = jax.random.PRNGKey(int(seed))
-    num_inference_steps = int(num_inference_steps)
-    image = Image.fromarray(image, mode="RGB")
-    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)
-    negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples)
-    processed_image = pipe.prepare_image_inputs([image] * num_samples)
-    
-    prompt_ids = shard(prompt_ids)
-    negative_prompt_ids = shard(negative_prompt_ids)
-    processed_image = shard(processed_image)
-    
-    output = pipe(
-        prompt_ids=prompt_ids,
-        image=processed_image,
-        params=p_params,
-        prng_seed=p_rng,
-        num_inference_steps=num_inference_steps,
-        neg_prompt_ids=negative_prompt_ids,
-        jit=True,
-    ).images
-
-    output = output.reshape((num_samples,) + output.shape[-3:])
-
-    print(output.shape)
+    try:
+        rng = jax.random.PRNGKey(int(seed))
+        num_inference_steps = int(num_inference_steps)
+        image = Image.fromarray(image, mode="RGB")
+        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)
+        negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples)
+        processed_image = pipe.prepare_image_inputs([image] * num_samples)
+        
+        prompt_ids = shard(prompt_ids)
+        negative_prompt_ids = shard(negative_prompt_ids)
+        processed_image = shard(processed_image)
+        
+        output = pipe(
+            prompt_ids=prompt_ids,
+            image=processed_image,
+            params=p_params,
+            prng_seed=p_rng,
+            num_inference_steps=num_inference_steps,
+            neg_prompt_ids=negative_prompt_ids,
+            jit=True,
+        ).images
+
+        del negative_prompt_ids
+        del processed_image
+        del prompt_ids
+
+        output = output.reshape((num_samples,) + output.shape[-3:])
+        final_image = [np.array(x*255, dtype=np.uint8) for x in output]
+        print(output.shape)
+        del output
+        
+    except Exception as e:
+        print("Error: " + str(e))
+        final_image = [np.zeros((512, 512, 3), dtype=np.uint8)] * num_samples
+    finally:
+        gc.collect()
+        return final_image
     
-    final_image = [np.array(x*255, dtype=np.uint8) for x in output]
 
-    del output
-    
-    return final_image
+default_example = examples[2]
+
+cond_img = gr.Image(label="Input", shape=(512, 512), value=default_example[2])\
+                    .style(height=200)
+
+output = gr.Gallery(label="Generated images")\
+                    .style(height=200, rows=[2], columns=[1, 2], object_fit="contain")
+
+prompt = gr.Textbox(lines=1, label="Prompt", value=default_example[0])
+negative_prompt = gr.Textbox(lines=1, label="Negative Prompt", value=default_example[1])
+
 
 with gr.Blocks(css="h1 { text-align: center }") as demo:
-    # Title
-    gr.Markdown(title)
-    # Description
-    gr.Markdown(description)
+    with gr.Row():
+        with gr.Column():
+            # Title
+            gr.Markdown(title)
+            # Description
+            gr.Markdown(description)
+
+        with gr.Column():
+            # Examples
+            gr.Examples(examples=examples,
+                    inputs=[prompt, negative_prompt, cond_img],
+                    outputs=output,
+                    fn=infer)
 
     # Images
     with gr.Row(variant="panel"):
         with gr.Column(scale=2):
-            cond_img = gr.Image(label="Input")\
-                    .style(height=200)
+            cond_img.render()
         with gr.Column(scale=1):
-            output = gr.Gallery(label="Generated images")\
-                    .style(height=200, rows=[2], columns=[1, 2], object_fit="contain")
+            output.render()
         
     # Submit & Clear
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(lines=1, label="Prompt")
-            negative_prompt = gr.Textbox(lines=1, label="Negative Prompt")
+            prompt.render()
+            negative_prompt.render()
 
         with gr.Column():
             with gr.Accordion("Advanced options", open=False):
@@ -102,13 +130,6 @@ with gr.Blocks(css="h1 { text-align: center }") as demo:
             submit = gr.Button("Generate")
             # TODO: Download Button
 
-    # Examples
-    gr.Examples(examples=examples,
-                inputs=[prompt, negative_prompt, cond_img],
-                outputs=output,
-                fn=infer,
-                cache_examples=True)
-
     
     submit.click(infer, 
                  inputs=[prompt, negative_prompt, cond_img, num_steps, seed, num_samples],