covert to blocks, enable webcam

app.py

@@ -35,16 +35,22 @@ pipe = pipe.to("cuda")
 # Generator seed,
 generator = torch.manual_seed(0)
 
+
 def get_bounding_box(image):
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    face = face_detector(gray)[0]
+    faces = face_detector(gray)
+    if len(faces) == 0:
+        raise Exception("No face detected in image")
+    face = faces[0]
     bbox = [face.left(), face.top(), face.width(), face.height()]
     return bbox
 
+
 def get_landmarks(image, bbox):
     features = spiga_extractor.inference(image, [bbox])
     return features['landmarks'][0]
 
+
 def get_patch(landmarks, color='lime', closed=False):
     contour = landmarks
     ops = [Path.MOVETO] + [Path.LINETO]*(len(contour)-1)
@@ -56,10 +62,12 @@ def get_patch(landmarks, color='lime', closed=False):
     path = Path(contour, ops)
     return patches.PathPatch(path, facecolor=facecolor, edgecolor=color, lw=4)
 
+
 def conditioning_from_landmarks(landmarks, size=512):
     # Precisely control output image size
     dpi = 72
-    fig, ax = plt.subplots(1, figsize=[size/dpi, size/dpi], tight_layout={'pad':0})
+    fig, ax = plt.subplots(
+        1, figsize=[size/dpi, size/dpi], tight_layout={'pad': 0})
     fig.set_dpi(dpi)
 
     black = np.zeros((size, size, 3))
@@ -86,17 +94,16 @@ def conditioning_from_landmarks(landmarks, size=512):
     ax.add_patch(inner_lips)
 
     plt.axis('off')
-    
     fig.canvas.draw()
     buffer, (width, height) = fig.canvas.print_to_buffer()
     assert width == height
     assert width == size
-    
     buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4))
     buffer = buffer[:, :, 0:3]
     plt.close(fig)
     return PIL.Image.fromarray(buffer)
 
+
 def get_conditioning(image):
     # Steps: convert to BGR and then:
     # - Retrieve bounding box using `dlib`
@@ -109,34 +116,72 @@ def get_conditioning(image):
     bbox = get_bounding_box(image)
     landmarks = get_landmarks(image, bbox)
     spiga_seg = conditioning_from_landmarks(landmarks)
-    return spiga_seg 
-
-def generate_images(image, prompt):
-    conditioning = get_conditioning(image)
-    output = pipe(
-        prompt,
-        conditioning,
-        generator=generator,
-        num_images_per_prompt=3,
-        num_inference_steps=20,
-    )
-    return [conditioning] + output.images
-
-
-gr.Interface(
-    generate_images,
-    inputs=[
-        gr.Image(type="pil"),
-        gr.Textbox(
-            label="Enter your prompt",
-            max_lines=1,
-            placeholder="best quality, extremely detailed",
-        ),
-    ],
-    outputs=gr.Gallery().style(grid=[2], height="auto"),
-    title="Generate controlled outputs with ControlNet and Stable Diffusion. ",
-    description="This Space uses a custom visualization based on SPIGA face landmarks for conditioning.",
-    # "happy zombie" instead of "young woman" works great too :)
-    examples=[["pedro-512.jpg", "Highly detailed photograph of young woman smiling, with palm trees in the background"]],
-    allow_flagging=False,
-).launch(enable_queue=True)
+    return spiga_seg
+
+
+def generate_images(image, prompt, image_video=None):
+    if image is None and image_video is None:
+        raise gr.Error("Please provide an image")
+    if image_video is not None:
+        image = image_video
+    try:
+        conditioning = get_conditioning(image)
+        output = pipe(
+            prompt,
+            conditioning,
+            generator=generator,
+            num_images_per_prompt=3,
+            num_inference_steps=20,
+        )
+        return [conditioning] + output.images
+    except Exception as e:
+        raise gr.Error(str(e))
+
+
+def toggle(choice):
+    if choice == "webcam":
+        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)
+
+
+with gr.Blocks() as blocks:
+    gr.Markdown("""
+        ## Generate controlled outputs with ControlNet and Stable Diffusion.
+        This Space uses a custom visualization based on SPIGA face landmarks for conditioning.
+    """)
+    with gr.Row():
+        with gr.Column():
+            image_or_file_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)
+            prompt = gr.Textbox(
+                label="Enter your prompt",
+                max_lines=1,
+                placeholder="best quality, extremely detailed",
+            )
+            run_button = gr.Button("Generate")
+        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])
+    gr.Examples(fn=generate_images,
+                examples=[
+                    ["./examples/pedro-512.jpg",
+                        "Highly detailed photograph of young woman smiling, with palm trees in the background"],
+                    ["./examples/image1.jpg",
+                        "Highly detailed photograph of a scary clown"],
+                    ["./examples/image0.jpg",
+                        "Highly detailed photograph of Barack Obama"],
+                ],
+                inputs=[image_in_img, prompt],
+                outputs=[gallery],
+                cache_examples=True)
+
+blocks.launch()