Spaces:
Runtime error
Runtime error
File size: 4,758 Bytes
ed82520 9c37edb ed82520 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
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 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)
|