diffusers/controlnet-openpose

1from controlnet_aux import OpenposeDetector2from diffusers import StableDiffusionControlNetPipeline, ControlNetModel3from diffusers import UniPCMultistepScheduler4import gradio as gr5import torch6import base647from io import BytesIO8from PIL import Image9# live conditioning10canvas_html = "<pose-canvas id='canvas-root' style='display:flex;max-width: 500px;margin: 0 auto;'></pose-canvas>"11load_js = """12async () => {13  const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/pose-gradio.js"14  fetch(url)15    .then(res => res.text())16    .then(text => {17      const script = document.createElement('script');18      script.type = "module"19      script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));20      document.head.appendChild(script);21    });22}23"""24get_js_image = """25async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {26  const canvasEl = document.getElementById("canvas-root");27  const data = canvasEl? canvasEl._data : null;28  return [image_in_img, prompt, image_file_live_opt, data]29}30"""31 32# Constants33low_threshold = 10034high_threshold = 20035 36# Models37pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")38controlnet = ControlNetModel.from_pretrained(39    "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float1640)41pipe = StableDiffusionControlNetPipeline.from_pretrained(42    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float1643)44pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)45 46# This command loads the individual model components on GPU on-demand. So, we don't47# need to explicitly call pipe.to("cuda").48pipe.enable_model_cpu_offload()49 50# xformers51pipe.enable_xformers_memory_efficient_attention()52 53# Generator seed,54generator = torch.manual_seed(0)55 56 57def get_pose(image):58    return pose_model(image)59 60 61def generate_images(image, prompt, image_file_live_opt='file', live_conditioning=None):62    if image is None and 'image' not in live_conditioning:63        raise gr.Error("Please provide an image")64    try:65        if image_file_live_opt == 'file':66            pose = get_pose(image)67        elif image_file_live_opt == 'webcam':68            base64_img = live_conditioning['image']69            image_data = base64.b64decode(base64_img.split(',')[1])70            pose = Image.open(BytesIO(image_data)).convert(71                'RGB').resize((512, 512))72        output = pipe(73            prompt,74            pose,75            generator=generator,76            num_images_per_prompt=3,77            num_inference_steps=20,78        )79        all_outputs = []80        all_outputs.append(pose)81        for image in output.images:82            all_outputs.append(image)83        return all_outputs84    except Exception as e:85        raise gr.Error(str(e))86 87 88def toggle(choice):89    if choice == "file":90        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)91    elif choice == "webcam":92        return gr.update(visible=False, value=None), gr.update(visible=True, value=canvas_html)93 94 95with gr.Blocks() as blocks:96    gr.Markdown("""97    ## Generate controlled outputs with ControlNet and Stable Diffusion98    This Space uses pose estimated lines as the additional conditioning99    [Check out our blog to see how this was done (and train your own controlnet)](https://huggingface.co/blog/train-your-controlnet)100    """)101    with gr.Row():102        live_conditioning = gr.JSON(value={}, visible=False)103        with gr.Column():104            image_file_live_opt = gr.Radio(["file", "webcam"], value="file",105                                           label="How would you like to upload your image?")106            image_in_img = gr.Image(source="upload", visible=True, type="pil")107            canvas = gr.HTML(None, elem_id="canvas_html", visible=False)108 109            image_file_live_opt.change(fn=toggle,110                                       inputs=[image_file_live_opt],111                                       outputs=[image_in_img, canvas],112                                       queue=False)113            prompt = gr.Textbox(114                label="Enter your prompt",115                max_lines=1,116                placeholder="best quality, extremely detailed",117            )118            run_button = gr.Button("Generate")119        with gr.Column():120            gallery = gr.Gallery().style(grid=[2], height="auto")121    run_button.click(fn=generate_images,122                     inputs=[image_in_img, prompt,123                             image_file_live_opt, live_conditioning],124                     outputs=[gallery],125                     _js=get_js_image)126    blocks.load(None, None, None, _js=load_js)127 128    gr.Examples(fn=generate_images,129                examples=[130                    ["./yoga1.jpeg",131                        "best quality, extremely detailed"]132                ],133                inputs=[image_in_img, prompt],134                outputs=[gallery],135                cache_examples=True)136 137blocks.launch(debug=True)138