ControlNet Head Mesh SDXL

• Paper: https://arxiv.org/abs/2407.18245

ControlNet Example(Conditioned on 3D Head Mesh)

👁 images_0)

Head Mesh Processor Install

pip install git+https://github.com/KupynOrest/head_detector.git

Code to Use Mesh Control

from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
from diffusers import EulerAncestralDiscreteScheduler
from PIL import Image
import torch
import numpy as np
import cv2

from head_detector import HeadDetector

detector = HeadDetector()

def resize_image(image):
 height, width = image.shape[:2]
 
 scale_factor = 1024.0 / max(height, width)
 new_width = int(width * scale_factor)
 new_height = int(height * scale_factor)

 resized_image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)

 return resized_image

def read_rgb_image(image_path):
 image = cv2.imread(image_path)
 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

 return image

IMAGE_PATH = "your image path"

controlnet_conditioning_scale = 1.0 
prompt = "your prompt, the longer the better, you can describe it as detail as possible"
negative_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'

eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")


controlnet = ControlNetModel.from_pretrained(
 "okupyn/head-mesh-controlnet-xl",
 torch_dtype=torch.float16
)

# when test with other base model, you need to change the vae also.
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
 "stabilityai/stable-diffusion-xl-base-1.0",
 controlnet=controlnet,
 vae=vae,
 safety_checker=None,
 torch_dtype=torch.float16,
 scheduler=eulera_scheduler,
)
#pipe = pipe.to("cuda")

image = resize_image(read_rgb_image(IMAGE_PATH))

controlnet_img = detector(image).get_pncc()
controlnet_img = Image.fromarray(controlnet_img)


images = pipe(
 prompt,
 negative_prompt=negative_prompt,
 image=controlnet_img,
 controlnet_conditioning_scale=controlnet_conditioning_scale,
 num_inference_steps=30,
 ).images

images[0].save(f"your image save path")

license: cc-by-nc-4.0 library_name: diffusers