do more of what canny_coyo1m does

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+import jax.numpy as jnp
+from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
 from diffusers import UniPCMultistepScheduler
 import torch
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -24,12 +25,12 @@ conditioning_image_transforms = T.Compose(
     ]
 )
 
-cnet = ControlNetModel.from_pretrained("./models/catcon-controlnet-wd", torch_dtype=torch.float16, from_flax=True).to("cuda")
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
+cnet = FlaxControlNetModel.from_pretrained("./models/catcon-controlnet-wd", dtype=jnp.bfloat16, from_flax=True)
+pipe = FlaxStableDiffusionControlNetPipeline.from_pretrained(
         "./models/wd-1-5-b2",
         controlnet=cnet,
-        torch_dtype=torch.float16,
-        ).to("cuda")
+        dtype=jnp.bfloat16,
+        )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 #pipe.enable_model_cpu_offload()
 #pipe.enable_xformers_memory_efficient_attention()
@@ -41,15 +42,18 @@ def infer(prompt, negative_prompt, image):
     # implement your inference function here
     inp = Image.fromarray(image)
 
-    cond_input = conditioning_image_transforms(inp).to("cpu", dtype=torch.float32)
+    cond_input = conditioning_image_transforms(inp)
     cond_input = T.ToPILImage()(cond_input)
+
+    cond_in = pipe.prepare_image_inputs([cond_input] * 4)
     
     output = pipe(
         prompt,
-        cond_input,
+        cond_in,
         generator=generator,
         num_images_per_prompt=4,
-        num_inference_steps=20
+        num_inference_steps=20,
+        jit=True
             )
 
     return output.images