tale/controlnet-3d-pose
0
1from diffusers import StableDiffusionControlNetPipeline, ControlNetModel2from diffusers import UniPCMultistepScheduler3import gradio as gr4import numpy as np5import torch6import base647import cv28from io import BytesIO9from PIL import Image, ImageFilter10 11from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css12 13# Constants14low_threshold = 10015high_threshold = 20016 17canvas_html = '<pose-maker/>'18load_js = """19async () => {20 const url = "https://huggingface.co/datasets/mishig/gradio-components/raw/main/mannequinAll.js"21 fetch(url)22 .then(res => res.text())23 .then(text => {24 const script = document.createElement('script');25 script.type = "module"26 script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));27 document.head.appendChild(script);28 });29}30"""31 32get_js_image = """33async (canvas, prompt) => {34 const poseMakerEl = document.querySelector("pose-maker");35 const imgBase64 = poseMakerEl.captureScreenshotDepthMap();36 return [imgBase64, prompt]37}38"""39 40# Models41controlnet = ControlNetModel.from_pretrained(42 "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float1643)44pipe = StableDiffusionControlNetPipeline.from_pretrained(45 "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float1646)47pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)48 49# This command loads the individual model components on GPU on-demand. So, we don't50# need to explicitly call pipe.to("cuda").51pipe.enable_model_cpu_offload()52 53# xformers54pipe.enable_xformers_memory_efficient_attention()55 56# Generator seed,57generator = torch.manual_seed(0)58 59def get_canny_filter(image):60 if not isinstance(image, np.ndarray):61 image = np.array(image) 62 63 image = cv2.Canny(image, low_threshold, high_threshold)64 image = image[:, :, None]65 image = np.concatenate([image, image, image], axis=2)66 canny_image = Image.fromarray(image)67 return canny_image68 69def generate_images(canvas, prompt):70 try:71 base64_img = canvas72 image_data = base64.b64decode(base64_img.split(',')[1])73 input_img = Image.open(BytesIO(image_data)).convert(74 'RGB').resize((512, 512))75 input_img = input_img.filter(ImageFilter.GaussianBlur(radius=2))76 input_img = get_canny_filter(input_img)77 output = pipe(78 f'{prompt}, unreal engine, Flickr, Canon camera, f50, best quality, extremely detailed',79 input_img,80 generator=generator,81 num_images_per_prompt=2,82 num_inference_steps=20,83 negative_prompt="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",84 )85 all_outputs = []86 for image in output.images:87 all_outputs.append(image)88 return all_outputs89 except Exception as e:90 raise gr.Error(str(e))91 92def placeholder_fn(axis):93 pass94 95js_change_rotation_axis = """96async (axis) => {97 const poseMakerEl = document.querySelector("pose-maker");98 poseMakerEl.changeRotationAxis(axis);99}100"""101 102js_pose_template = """103async (pose) => {104 const poseMakerEl = document.querySelector("pose-maker");105 poseMakerEl.setPose(pose);106}107"""108 109with gr.Blocks(css=share_btn_css) as blocks:110 gr.HTML(111 """112 <div style="text-align: center; margin: 0 auto;">113 <div114 style="115 display: inline-flex;116 align-items: center;117 gap: 0.8rem;118 font-size: 1.75rem;119 "120 >121 <h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">122 Pose in 3D & Render with ControlNet (SD-1.5)123 </h1>124 </div>125 <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">126 Using <a href="https://huggingface.co/blog/controlnet">ControlNet</a> and <a href="https://boytchev.github.io/mannequin.js/">three.js/mannequin.js</a>127 </p>128 <p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/diffusers/controlnet-3d-pose?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>129 </div>130 """131 )132 with gr.Row():133 with gr.Column():134 canvas = gr.HTML(canvas_html, elem_id="canvas_html", visible=True)135 with gr.Row():136 rotation_axis = gr.Radio(["x", "y", "z"], value="x", label="Joint rotation axis")137 pose_template = gr.Radio(["regular", "ballet", "handstand", "split", "kick", "chilling"], value="regular", label="Pose template")138 prompt = gr.Textbox(139 label="Enter your prompt",140 max_lines=1,141 placeholder="best quality, extremely detailed",142 elem_id="prompt",143 )144 run_button = gr.Button("Generate")145 gr.Markdown("### See an example [here](https://huggingface.co/spaces/diffusers/controlnet-3d-pose/discussions/1)")146 with gr.Group(elem_id="share-btn-container"):147 community_icon = gr.HTML(community_icon_html)148 loading_icon = gr.HTML(loading_icon_html)149 share_button = gr.Button("Share to community", elem_id="share-btn")150 with gr.Column():151 gallery = gr.Gallery(elem_id="gallery").style(grid=[2], height="auto")152 rotation_axis.change(fn=placeholder_fn,153 inputs=[rotation_axis],154 outputs=[],155 queue=False,156 _js=js_change_rotation_axis)157 pose_template.change(fn=placeholder_fn,158 inputs=[pose_template],159 outputs=[],160 queue=False,161 _js=js_pose_template)162 run_button.click(fn=generate_images,163 inputs=[canvas, prompt],164 outputs=[gallery],165 _js=get_js_image)166 share_button.click(None, [], [], _js=share_js)167 blocks.load(None, None, None, _js=load_js)168 169blocks.launch(debug=True)170