Update app.py

Adding control image

app.py

@@ -18,24 +18,24 @@ pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
 )
 
-#def infer(prompt, image):
-def infer(prompt):
+def infer(prompt, image):
+#def infer(prompt):
     params["controlnet"] = controlnet_params
     
     num_samples = 1 #jax.device_count()
     rng = create_key(0)
     rng = jax.random.split(rng, jax.device_count())
-    #im = image
-    #image = Image.fromarray(im)
+    im = image
+    image = Image.fromarray(im)
     
     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)
+    processed_image = pipe.prepare_image_inputs([image] * num_samples)
     
     p_params = replicate(params)
     prompt_ids = shard(prompt_ids)
     #negative_prompt_ids = shard(negative_prompt_ids)
-    #processed_image = shard(processed_image)
+    processed_image = shard(processed_image)
     
     output = pipe(
         prompt_ids=prompt_ids,
@@ -53,13 +53,6 @@ def infer(prompt):
 #gr.Interface(infer, inputs=["text", "text", "image"], outputs="gallery").launch()
 
 
-
-
-
-
-
-
-
 title = "Animal Pose Control Net"
 description = "This is a demo of Animal Pose ControlNet, which is a model trained on runwayml/stable-diffusion-v1-5 with new type of conditioning."
 
@@ -75,14 +68,13 @@ description = "This is a demo of Animal Pose ControlNet, which is a model traine
 #    button_primary_background_fill_hover="*primary_300",
 #)
 
-#gr.Interface(fn = infer, inputs = ["text", "text", "image"], outputs = "image",
-gr.Interface(fn = infer, inputs = ["text"], outputs = "image",
-            title = title, description = description, theme='gradio/soft').launch()
+#gr.Interface(fn = infer, inputs = ["text"], outputs = "image",
+#            title = title, description = description, theme='gradio/soft').launch()
 
-#gr.Interface(fn = infer, inputs = ["text", "text", "image"], outputs = "gallery",
- #            title = title, description = description, theme='gradio/soft',
-  #           examples=[["a Labrador crossing the road", "low quality", "pose_256.jpg"]]
-#).launch()
+gr.Interface(fn = infer, inputs = ["text", "text", "image"], outputs = "gallery",
+            title = title, description = description, theme='gradio/soft',
+            examples=[["a Labrador crossing the road", "low quality", "image_control.png"]]
+).launch()
 
 gr.Markdown(
         """