Update app.py

Using https://huggingface.co/spaces/jax-diffusers-event/canny_coyo1m/blob/main/app.py as a guide

app.py

@@ -1,60 +1,58 @@
-from PIL import Image
 import gradio as gr
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
-import torch
-torch.backends.cuda.matmul.allow_tf32 = True
+import jax.numpy as jnp
+import jax
+import numpy as np
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+from PIL import Image
+from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
+import cv2
 
-controlnet = ControlNetModel.from_pretrained("JFoz/dog-cat-pose", torch_dtype=torch.float16)
+def create_key(seed=0):
+    return jax.random.PRNGKey(seed)
 
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    controlnet=controlnet,
-    torch_dtype=torch.float16,
-    safety_checker=None,
+controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
+    "JFoz/dog-cat-pose", dtype=jnp.bfloat16
+)
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
 )
 
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-
-pipe.enable_xformers_memory_efficient_attention()
-pipe.enable_model_cpu_offload()
-pipe.enable_attention_slicing()
-
-def infer(
-        prompt,
-        negative_prompt,
-        conditioning_image,
-        num_inference_steps=30,
-        size=768,
-        guidance_scale=7.0,
-        seed=1234,
-):
-
-    conditioning_image_raw = Image.fromarray(conditioning_image)
-    #conditioning_image = conditioning_image_raw.convert('L')
-
-    g_cpu = torch.Generator()
-
-    if seed == -1:
-        generator = g_cpu.manual_seed(g_cpu.seed())
-    else:
-        generator = g_cpu.manual_seed(seed)
-
-    output_image = pipe(
-        prompt,
-        conditioning_image,
-        height=size,
-        width=size,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        controlnet_conditioning_scale=1.0,
-    ).images[0]
+def infer(prompts, negative_prompts, image):
+    params["controlnet"] = controlnet_params
+    
+    num_samples = 1 #jax.device_count()
+    rng = create_key(0)
+    rng = jax.random.split(rng, jax.device_count())
+    #im = canny_filter(image)
+    #canny_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([canny_image] * num_samples)
+    
+    p_params = replicate(params)
+    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=rng,
+        num_inference_steps=50,
+        neg_prompt_ids=negative_prompt_ids,
+        jit=True,
+    ).images
+    
+    output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
+    return output_images
 
-    #del conditioning_image, conditioning_image_raw
-    #gc.collect()
+#gr.Interface(infer, inputs=["text", "text", "image"], outputs="gallery").launch()
 
-    return output_image
+title = "Animal Pose Control Net"
+description = "This is a demo on ControlNet based on canny filter."
 
 with gr.Blocks(theme=gr.themes.Default(font=[gr.themes.GoogleFont("Inconsolata"), "Arial", "sans-serif"])) as demo:
     gr.Markdown(
@@ -63,83 +61,22 @@ with gr.Blocks(theme=gr.themes.Default(font=[gr.themes.GoogleFont("Inconsolata")
     # This is a demo of Animal Pose Control Net, which is a model trained on runwayml/stable-diffusion-v1-5 with new type of conditioning.
     """)
 
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(
-                label="Prompt",
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-            )
-            conditioning_image = gr.Image(
-                label="Conditioning Image",
-            )
-            with gr.Accordion('Advanced options', open=False):
-                with gr.Row():
-                    num_inference_steps = gr.Slider(
-                        10, 40, 20,
-                        step=1,
-                        label="Steps",
-                    )
-                    size = gr.Slider(
-                        256, 768, 512,
-                        step=128,
-                        label="Size",
-                    )
-                with gr.Row():
-                    guidance_scale = gr.Slider(
-                        label='Guidance Scale',
-                        minimum=0.1,
-                        maximum=30.0,
-                        value=7.0,
-                        step=0.1
-                    )
-                    seed = gr.Slider(
-                        label='Seed',
-                        value=-1,
-                        minimum=-1,
-                        maximum=2147483647,
-                        step=1,
-                        # randomize=True
-                    )
-            submit_btn = gr.Button(
-                value="Submit",
-                variant="primary"
-            )
-        with gr.Column(min_width=300):
-            output = gr.Image(
-                label="Result",
-            )
-
-    submit_btn.click(
-        fn=infer,
-        inputs=[
-            prompt, negative_prompt, conditioning_image, num_inference_steps, size, guidance_scale, seed
-            #prompt, size, seed
-        ],
-        outputs=output
-    )
-    gr.Examples(
+gr.Examples(
         examples=[
             #["a tortoiseshell cat is sitting on a cushion"],
             #["a yellow dog standing on a lawn"],
             ["a tortoiseshell cat is sitting on a cushion",  "https://huggingface.co/JFoz/dog-cat-pose/blob/main/images_0.png"],
             ["a yellow dog standing on a lawn", "https://huggingface.co/JFoz/dog-cat-pose/blob/main/images_1.png"],
-        ],
-        inputs=[
-            #prompt, negative_prompt, conditioning_image
-            prompt
-        ],
-        outputs=output,
-        fn=infer,
+        ]
         cache_examples=True,
     )
-    gr.Markdown(
+
+gr.Interface(fn = infer, inputs = ["text", "text", "image"], outputs = "image",
+            title = title, description = description, examples = gr.examples, theme='gradio/soft').launch()
+
+gr.Markdown(
         """
     * [Dataset](https://huggingface.co/datasets/JFoz/dog-poses-controlnet-dataset)
     * [Diffusers model](), [Web UI model](https://huggingface.co/JFoz/dog-pose)
     * [Training Report](https://wandb.ai/john-fozard/dog-cat-pose/runs/kmwcvae5))
     """)
-
-#gr.Interface(infer, inputs=["text"], outputs=[output], title=title, description=description, examples=examples).queue().launch()
-demo.launch()