.

app.py

@@ -33,8 +33,7 @@ pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
         "./models/wd-1-5-b2-flax",
         controlnet=cnet,
         revision="flax",
-        dtype=jnp.bfloat16,
-        safety_checker=None,
+        dtype=jnp.bfloat16
         )
 #pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 #pipe.enable_model_cpu_offload()
@@ -52,7 +51,7 @@ def infer(prompt, negative_prompt, image):
     inp = Image.fromarray(image)
 
     cond_input = conditioning_image_transforms(inp)
-    #cond_input = T.ToPILImage()(cond_input)
+    cond_input = T.ToPILImage()(cond_input)
 
     cond_img_in = pipe.prepare_image_inputs([cond_input] * num_samples)
     cond_img_in = shard(cond_img_in)