enable live conditioning (#2)

- enable live conditioning (219f54ab9f753d37099cb613ad68703220b44db9)


Co-authored-by: Radamés Ajna <[email protected]>

app.py

@@ -3,7 +3,9 @@ import torch
 import dlib
 import numpy as np
 import PIL
-
+import base64
+from io import BytesIO
+from PIL import Image
 # Only used to convert to gray, could do it differently and remove this big dependency
 import cv2
 
@@ -35,6 +37,26 @@ pipe = pipe.to("cuda")
 # Generator seed,
 generator = torch.manual_seed(0)
 
+canvas_html = "<face-canvas id='canvas-root' style='display:flex;max-width: 500px;margin: 0 auto;'></face-canvas>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/face-canvas.js"
+fetch(url)
+  .then(res => res.text())
+  .then(text => {
+    const script = document.createElement('script');
+    script.type = "module"
+    script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
+    document.head.appendChild(script);
+  });
+}
+"""
+get_js_image = """
+async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [image_in_img, prompt, image_file_live_opt, canvasEl._data]
+}
+"""
 
 def get_bounding_box(image):
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
@@ -119,13 +141,18 @@ def get_conditioning(image):
     return spiga_seg
 
 
-def generate_images(image, prompt, image_video=None):
-    if image is None and image_video is None:
+def generate_images(image_in_img, prompt, image_file_live_opt='file', live_conditioning=None):
+    if image_in_img is None and 'image' not in live_conditioning:
         raise gr.Error("Please provide an image")
-    if image_video is not None:
-        image = image_video
+
+    if image_file_live_opt == 'file':
+      conditioning = get_conditioning(image_in_img)
+    elif image_file_live_opt == 'webcam':
+      base64_img = live_conditioning['image']
+      image_data = base64.b64decode(base64_img.split(',')[1])
+      conditioning = Image.open(BytesIO(image_data)).convert('RGB').resize((512,512))
+
     try:
-        conditioning = get_conditioning(image)
         output = pipe(
             prompt,
             conditioning,
@@ -139,11 +166,10 @@ def generate_images(image, prompt, image_video=None):
 
 
 def toggle(choice):
-    if choice == "webcam":
+    if choice == "file":
         return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
-    else:
-        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
-
+    elif choice == "webcam":
+        return gr.update(visible=False, value=None), gr.update(visible=True, value=canvas_html)
 
 with gr.Blocks() as blocks:
     gr.Markdown("""
@@ -151,15 +177,17 @@ with gr.Blocks() as blocks:
         [Check out our blog to see how this was done (and train your own controlnet)](https://huggingface.co/blog/train-your-controlnet)
     """)
     with gr.Row():
+        live_conditioning  = gr.JSON(value={}, visible=False)
         with gr.Column():
-            image_or_file_opt = gr.Radio(["file", "webcam"], value="file",
+            image_file_live_opt = gr.Radio(["file", "webcam"], value="file",
                                          label="How would you like to upload your image?")
-            image_in_video = gr.Image(
-                source="webcam", type="pil", visible=False)
-            image_in_img = gr.Image(
-                source="upload", visible=True, type="pil")
-            image_or_file_opt.change(fn=toggle, inputs=[image_or_file_opt],
-                                     outputs=[image_in_video, image_in_img], queue=False)
+            image_in_img = gr.Image(source="upload", visible=True, type="pil")
+            canvas = gr.HTML(None, elem_id="canvas_html", visible=False)
+
+            image_file_live_opt.change(fn=toggle,
+                                     inputs=[image_file_live_opt],
+                                     outputs=[image_in_img, canvas],
+                                     queue=False)
             prompt = gr.Textbox(
                 label="Enter your prompt",
                 max_lines=1,
@@ -169,8 +197,10 @@ with gr.Blocks() as blocks:
         with gr.Column():
             gallery = gr.Gallery().style(grid=[2], height="auto")
     run_button.click(fn=generate_images,
-                     inputs=[image_in_img, prompt, image_in_video],
-                     outputs=[gallery])
+                     inputs=[image_in_img, prompt, image_file_live_opt, live_conditioning],
+                     outputs=[gallery],
+                     _js=get_js_image)
+    blocks.load(None, None, None, _js=load_js)
     gr.Examples(fn=generate_images,
                 examples=[
                     ["./examples/pedro-512.jpg",
@@ -178,7 +208,7 @@ with gr.Blocks() as blocks:
                     ["./examples/image1.jpg",
                         "Highly detailed photograph of a scary clown"],
                     ["./examples/image0.jpg",
-                        "Highly detailed photograph of Barack Obama"],
+                        "Highly detailed photograph of Madonna"],
                 ],
                 inputs=[image_in_img, prompt],
                 outputs=[gallery],