Initial version

README.md

@@ -1,10 +1,10 @@
 ---
-title: Uncanny Faces
-emoji: 🏢
-colorFrom: pink
-colorTo: pink
+title: ControlNet Openpose
+emoji: 😻
+colorFrom: green
+colorTo: gray
 sdk: gradio
-sdk_version: 3.23.0
+sdk_version: 3.19.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,142 @@
+import gradio as gr
+import torch
+import dlib
+import numpy as np
+import PIL
+
+# Only used to convert to gray, could do it differently and remove this big dependency
+import cv2
+
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import UniPCMultistepScheduler
+
+from spiga.inference.config import ModelConfig
+from spiga.inference.framework import SPIGAFramework
+
+import matplotlib.pyplot as plt
+from matplotlib.path import Path
+import matplotlib.patches as patches
+
+# Bounding boxes
+face_detector = dlib.get_frontal_face_detector()
+
+# Landmark extraction
+spiga_extractor = SPIGAFramework(ModelConfig("300wpublic"))
+
+uncanny_controlnet = ControlNetModel.from_pretrained(
+    "multimodalart/uncannyfaces_25K", torch_dtype=torch.float16
+)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1-base", controlnet=uncanny_controlnet, safety_checker=None, torch_dtype=torch.float16
+)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+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]
+    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)
+    facecolor = (0, 0, 0, 0)      # Transparent fill color, if open
+    if closed:
+        contour.append(contour[0])
+        ops.append(Path.CLOSEPOLY)
+        facecolor = color
+    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.set_dpi(dpi)
+
+    black = np.zeros((size, size, 3))
+    ax.imshow(black)
+
+    face_patch = get_patch(landmarks[0:17])
+    l_eyebrow = get_patch(landmarks[17:22], color='yellow')
+    r_eyebrow = get_patch(landmarks[22:27], color='yellow')
+    nose_v = get_patch(landmarks[27:31], color='orange')
+    nose_h = get_patch(landmarks[31:36], color='orange')
+    l_eye = get_patch(landmarks[36:42], color='magenta', closed=True)
+    r_eye = get_patch(landmarks[42:48], color='magenta', closed=True)
+    outer_lips = get_patch(landmarks[48:60], color='cyan', closed=True)
+    inner_lips = get_patch(landmarks[60:68], color='blue', closed=True)
+
+    ax.add_patch(face_patch)
+    ax.add_patch(l_eyebrow)
+    ax.add_patch(r_eyebrow)
+    ax.add_patch(nose_v)
+    ax.add_patch(nose_h)
+    ax.add_patch(l_eye)
+    ax.add_patch(r_eye)
+    ax.add_patch(outer_lips)
+    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`
+    # - Obtain landmarks using `spiga`
+    # - Create conditioning image with custom `matplotlib` code
+    # TODO: error if bbox is too small
+    image.thumbnail((512, 512))
+    image = np.array(image)
+    image = image[:, :, ::-1]
+    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 pose estimated lines as the additional 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)

requirements.txt

@@ -0,0 +1,9 @@
+diffusers
+transformers
+accelerate
+torch
+git+https://github.com/andresprados/SPIGA
+dlib
+opencv-python
+matplotlib
+Pillow