hysts/ControlNet
993
1#!/usr/bin/env python2 3from __future__ import annotations4 5import os6import pathlib7import shlex8import subprocess9 10import gradio as gr11import torch12 13if os.getenv('SYSTEM') == 'spaces':14 with open('patch') as f:15 subprocess.run(shlex.split('patch -p1'), stdin=f, cwd='ControlNet')16 17base_url = 'https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/'18names = [19 'body_pose_model.pth',20 'dpt_hybrid-midas-501f0c75.pt',21 'hand_pose_model.pth',22 'mlsd_large_512_fp32.pth',23 'mlsd_tiny_512_fp32.pth',24 'network-bsds500.pth',25 'upernet_global_small.pth',26]27for name in names:28 command = f'wget https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/{name} -O {name}'29 out_path = pathlib.Path(f'ControlNet/annotator/ckpts/{name}')30 if out_path.exists():31 continue32 subprocess.run(shlex.split(command), cwd='ControlNet/annotator/ckpts/')33 34from app_canny import create_demo as create_demo_canny35from app_depth import create_demo as create_demo_depth36from app_fake_scribble import create_demo as create_demo_fake_scribble37from app_hed import create_demo as create_demo_hed38from app_hough import create_demo as create_demo_hough39from app_normal import create_demo as create_demo_normal40from app_pose import create_demo as create_demo_pose41from app_scribble import create_demo as create_demo_scribble42from app_scribble_interactive import \43 create_demo as create_demo_scribble_interactive44from app_seg import create_demo as create_demo_seg45from model import Model, download_all_controlnet_weights46 47DESCRIPTION = '''# [ControlNet v1.0](https://github.com/lllyasviel/ControlNet)48 49<p class="note">New ControlNet v1.1 is available <a href="https://huggingface.co/spaces/hysts/ControlNet-v1-1">here</a>.</p>50'''51 52SPACE_ID = os.getenv('SPACE_ID')53ALLOW_CHANGING_BASE_MODEL = SPACE_ID != 'hysts/ControlNet'54 55if SPACE_ID is not None:56 DESCRIPTION += f'\n<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/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'57if not torch.cuda.is_available():58 DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'59 60if torch.cuda.is_available():61 if os.getenv('SYSTEM') == 'spaces':62 download_all_controlnet_weights()63 64MAX_IMAGES = int(os.getenv('MAX_IMAGES', '3'))65DEFAULT_NUM_IMAGES = min(MAX_IMAGES, int(os.getenv('DEFAULT_NUM_IMAGES', '1')))66 67DEFAULT_MODEL_ID = os.getenv('DEFAULT_MODEL_ID',68 'runwayml/stable-diffusion-v1-5')69model = Model(base_model_id=DEFAULT_MODEL_ID, task_name='canny')70 71with gr.Blocks(css='style.css') as demo:72 gr.Markdown(DESCRIPTION)73 with gr.Tabs():74 with gr.TabItem('Canny'):75 create_demo_canny(model.process_canny,76 max_images=MAX_IMAGES,77 default_num_images=DEFAULT_NUM_IMAGES)78 with gr.TabItem('Hough'):79 create_demo_hough(model.process_hough,80 max_images=MAX_IMAGES,81 default_num_images=DEFAULT_NUM_IMAGES)82 with gr.TabItem('HED'):83 create_demo_hed(model.process_hed,84 max_images=MAX_IMAGES,85 default_num_images=DEFAULT_NUM_IMAGES)86 with gr.TabItem('Scribble'):87 create_demo_scribble(model.process_scribble,88 max_images=MAX_IMAGES,89 default_num_images=DEFAULT_NUM_IMAGES)90 with gr.TabItem('Scribble Interactive'):91 create_demo_scribble_interactive(92 model.process_scribble_interactive,93 max_images=MAX_IMAGES,94 default_num_images=DEFAULT_NUM_IMAGES)95 with gr.TabItem('Fake Scribble'):96 create_demo_fake_scribble(model.process_fake_scribble,97 max_images=MAX_IMAGES,98 default_num_images=DEFAULT_NUM_IMAGES)99 with gr.TabItem('Pose'):100 create_demo_pose(model.process_pose,101 max_images=MAX_IMAGES,102 default_num_images=DEFAULT_NUM_IMAGES)103 with gr.TabItem('Segmentation'):104 create_demo_seg(model.process_seg,105 max_images=MAX_IMAGES,106 default_num_images=DEFAULT_NUM_IMAGES)107 with gr.TabItem('Depth'):108 create_demo_depth(model.process_depth,109 max_images=MAX_IMAGES,110 default_num_images=DEFAULT_NUM_IMAGES)111 with gr.TabItem('Normal map'):112 create_demo_normal(model.process_normal,113 max_images=MAX_IMAGES,114 default_num_images=DEFAULT_NUM_IMAGES)115 116 with gr.Accordion(label='Base model', open=False):117 with gr.Row():118 with gr.Column():119 current_base_model = gr.Text(label='Current base model')120 with gr.Column(scale=0.3):121 check_base_model_button = gr.Button('Check current base model')122 with gr.Row():123 with gr.Column():124 new_base_model_id = gr.Text(125 label='New base model',126 max_lines=1,127 placeholder='runwayml/stable-diffusion-v1-5',128 info=129 'The base model must be compatible with Stable Diffusion v1.5.',130 interactive=ALLOW_CHANGING_BASE_MODEL)131 with gr.Column(scale=0.3):132 change_base_model_button = gr.Button(133 'Change base model', interactive=ALLOW_CHANGING_BASE_MODEL)134 if not ALLOW_CHANGING_BASE_MODEL:135 gr.Markdown(136 '''The base model is not allowed to be changed in this Space so as not to slow down the demo, but it can be changed if you duplicate the Space. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a>'''137 )138 139 gr.Markdown('''### Related Spaces140 141- [Space using Anything-v4.0 as base model](https://huggingface.co/spaces/hysts/ControlNet-with-Anything-v4)142- https://huggingface.co/spaces/jonigata/PoseMaker2143- https://huggingface.co/spaces/diffusers/controlnet-openpose144- https://huggingface.co/spaces/diffusers/controlnet-canny145''')146 147 check_base_model_button.click(fn=lambda: model.base_model_id,148 outputs=current_base_model,149 queue=False)150 new_base_model_id.submit(fn=model.set_base_model,151 inputs=new_base_model_id,152 outputs=current_base_model)153 change_base_model_button.click(fn=model.set_base_model,154 inputs=new_base_model_id,155 outputs=current_base_model)156 157demo.queue(api_open=False, max_size=10).launch()158