use revision=flax

app.py

@@ -31,8 +31,8 @@ cnet, cnet_params = FlaxControlNetModel.from_pretrained("./models/catcon-control
 pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
         "./models/wd-1-5-b2",
         controlnet=cnet,
+        revision="flax"
         dtype=jnp.bfloat16,
-        from_pt=True
         )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 #pipe.enable_model_cpu_offload()
@@ -72,9 +72,10 @@ def infer(prompt, negative_prompt, image):
         neg_prompt_ids=n_prompt_in,
         num_inference_steps=20,
         jit=True
-            )
+            ).images
 
-    return output.images
+    output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
+    return output_images
 
 gr.Interface(
     infer,