import gradio as gr import torch import numpy as np import PIL import base64 from io import BytesIO from PIL import Image # import for face detection import retinaface from diffusers import StableDiffusionControlNetPipeline, ControlNetModel from diffusers import UniPCMultistepScheduler from spiga.inference.config import ModelConfig from spiga.inference.framework import SPIGAFramework import spiga.demo.analyze.track.retinasort.config as cfg import matplotlib.pyplot as plt from matplotlib.path import Path import matplotlib.patches as patches # Bounding boxes config = cfg.cfg_retinasort face_detector = retinaface.RetinaFaceDetector(model=config['retina']['model_name'], device='cuda' if torch.cuda.is_available() else 'cpu', extra_features=config['retina']['extra_features'], cfg_postreat=config['retina']['postreat']) # 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) canvas_html = "" 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"); const imageData = canvasEl? canvasEl._data : null; return [image_in_img, prompt, image_file_live_opt, imageData] } """ def get_bounding_box(image): pil_image = Image.fromarray(image) face_detector.set_input_shape(pil_image.size[1], pil_image.size[0]) features = face_detector.inference(pil_image) if (features is None) and (len(features['bbox']) <= 0): raise Exception("No face detected") # get the first face detected bbox = features['bbox'][0] x1, y1, x2, y2 = bbox[:4] bbox_wh = [x1, y1, x2-x1, y2-y1] return bbox_wh 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_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") try: 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)) 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 == "file": return gr.update(visible=True, value=None), gr.update(visible=False, 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(""" ## Generate Uncanny Faces with ControlNet Stable Diffusion [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_file_live_opt = gr.Radio(["file", "webcam"], value="file", label="How would you like to upload your image?") 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, 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_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", "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 Madonna"], ], inputs=[image_in_img, prompt], outputs=[gallery], cache_examples=True) gr.Markdown(''' This Space was trained on synthetic 3D faces to learn how to keep a pose - however it also learned that all faces are synthetic 3D faces, [learn more on our blog](https://huggingface.co/blog/train-your-controlnet), it uses a custom visualization based on SPIGA face landmarks for conditioning. ''') blocks.launch()