Update app.py

app.py

@@ -1,60 +1,48 @@
 import gradio as gr
-from datasets import load_dataset
 import jax
 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 diffusers.utils.testing_utils import load_image
 from PIL import Image
 from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
-
-def image_grid(imgs, rows, cols):
-    w, h = imgs[0].size
-    grid = Image.new("RGB", size=(cols * w, rows * h))
-    for i, img in enumerate(imgs):
-        grid.paste(img, box=(i % cols * w, i // cols * h))
-    return grid
+import cv2
 
 def create_key(seed=0):
     return jax.random.PRNGKey(seed)
 
-
-def infer(prompt, negative_prompt, image):
-    rng = create_key(0)
-
-#    canny_image = load_image(
-#        "https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/blog_post_cell_10_output_0.jpeg"
-#    )
-    canny_image = load_image(image)
-
-    #prompts = "a living room fan"
-    prompts = prompt
-    negative_prompts = negative_prompt
-
-    controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
-        "tsungtao/controlnet-mlsd-202305011046", from_flax=True, dtype=jnp.float32
-    )
-
-    pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.float32
-    )
-
+def canny_filter(image):
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
+    edges_image = cv2.Canny(blurred_image, 50, 150)
+    return edges_image
+
+# load control net and stable diffusion v1-5
+controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
+    "tsungtao/controlnet-mlsd-202305011046", from_flax=True, dtype=jnp.bfloat16
+)
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
+)
+
+def infer(prompts, negative_prompts, image):
     params["controlnet"] = controlnet_params
-
-    num_samples = jax.device_count()
+    
+    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,
@@ -64,25 +52,8 @@ def infer(prompt, negative_prompt, image):
         neg_prompt_ids=negative_prompt_ids,
         jit=True,
     ).images
-
+    
     output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
-    output_images = image_grid(output_images, num_samples // 4, 4)
-    #output_images.save("tao/image.png")
-    #dataset = load_dataset('imagefolder', data_dir='tao')
-    #dataset.push_to_hub('tsungtao/tmp')
     return output_images
 
-#infer('','','')
-
-def infer2(prompt, negative_prompt, image):
-    output_image = infer(prompt, negative_prompt, image)
-    #output_image = "https://datasets-server.huggingface.co/assets/tsungtao/tmp/--/tsungtao--tmp/train/0/image/image.jpg"
-    return output_image
-
-title = "ControlNet on MLSD Filter"
-description = "This is a demo on ControlNet based on mlsd filter."
-
-#examples = [["living room with TV", "fan", "https://datasets-server.huggingface.co/assets/tsungtao/diffusers-testing/--/tsungtao--diffusers-testing/train/0/images/image.jpg"]]
-
-interface = gr.Interface(fn = infer2, inputs = ["text", "text", "text"], outputs = "image",title = title, description = description, theme='gradio/soft')
-interface.launch(enable_queue=True)
+gr.Interface(infer, inputs=["text", "text", "image"], outputs="gallery").launch()