Update app.py

app.py

@@ -4,10 +4,25 @@ import numpy as np
 import jax.numpy as jnp
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
-from diffusers.utils import load_image
 from PIL import Image
 from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
 import cv2
+import os
+
+
+def load_image(image):
+    if isinstance(image, str):
+        if image.startswith("http://") or image.startswith("https://"):
+            image = PIL.Image.open(requests.get(image, stream=True).raw)
+        elif os.path.isfile(image):
+            image = PIL.Image.open(image)
+    elif isinstance(image, PIL.Image.Image):
+        image = image
+    image = PIL.ImageOps.exif_transpose(image)
+    image = image.convert("RGB")
+    return image
+
+
 
 def image_grid(imgs, rows, cols):
     w, h = imgs[0].size