zoro005/ControlNet
0
1#!/usr/bin/env python2 3from __future__ import annotations4 5import os6import pathlib7import shlex8import subprocess9 10import gradio as gr11 12if os.getenv('SYSTEM') == 'spaces':13 with open('patch') as f:14 subprocess.run(shlex.split('patch -p1'), stdin=f, cwd='ControlNet')15 16base_url = 'https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/'17names = [18 'body_pose_model.pth',19 'dpt_hybrid-midas-501f0c75.pt',20 'hand_pose_model.pth',21 'mlsd_large_512_fp32.pth',22 'mlsd_tiny_512_fp32.pth',23 'network-bsds500.pth',24 'upernet_global_small.pth',25]26for name in names:27 command = f'wget https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/{name} -O {name}'28 out_path = pathlib.Path(f'ControlNet/annotator/ckpts/{name}')29 if out_path.exists():30 continue31 subprocess.run(shlex.split(command), cwd='ControlNet/annotator/ckpts/')32 33from app_canny import create_demo as create_demo_canny34from app_depth import create_demo as create_demo_depth35from app_fake_scribble import create_demo as create_demo_fake_scribble36from app_hed import create_demo as create_demo_hed37from app_hough import create_demo as create_demo_hough38from app_normal import create_demo as create_demo_normal39from app_pose import create_demo as create_demo_pose40from app_scribble import create_demo as create_demo_scribble41from app_scribble_interactive import \42 create_demo as create_demo_scribble_interactive43from app_seg import create_demo as create_demo_seg44from model import Model, download_all_controlnet_weights45 46DESCRIPTION = '# [ControlNet](https://github.com/lllyasviel/ControlNet)'47 48SPACE_ID = os.getenv('SPACE_ID')49ALLOW_CHANGING_BASE_MODEL = SPACE_ID != 'hysts/ControlNet'50 51if SPACE_ID is not None:52 DESCRIPTION += f'<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>'53 54MAX_IMAGES = int(os.getenv('MAX_IMAGES', '3'))55DEFAULT_NUM_IMAGES = min(MAX_IMAGES, int(os.getenv('DEFAULT_NUM_IMAGES', '1')))56 57if os.getenv('SYSTEM') == 'spaces':58 download_all_controlnet_weights()59 60DEFAULT_MODEL_ID = os.getenv('DEFAULT_MODEL_ID',61 'runwayml/stable-diffusion-v1-5')62model = Model(base_model_id=DEFAULT_MODEL_ID, task_name='canny')63 64with gr.Blocks(css='style.css') as demo:65 gr.Markdown(DESCRIPTION)66 with gr.Tabs():67 with gr.TabItem('Canny'):68 create_demo_canny(model.process_canny,69 max_images=MAX_IMAGES,70 default_num_images=DEFAULT_NUM_IMAGES)71 with gr.TabItem('Hough'):72 create_demo_hough(model.process_hough,73 max_images=MAX_IMAGES,74 default_num_images=DEFAULT_NUM_IMAGES)75 with gr.TabItem('HED'):76 create_demo_hed(model.process_hed,77 max_images=MAX_IMAGES,78 default_num_images=DEFAULT_NUM_IMAGES)79 with gr.TabItem('Scribble'):80 create_demo_scribble(model.process_scribble,81 max_images=MAX_IMAGES,82 default_num_images=DEFAULT_NUM_IMAGES)83 with gr.TabItem('Scribble Interactive'):84 create_demo_scribble_interactive(85 model.process_scribble_interactive,86 max_images=MAX_IMAGES,87 default_num_images=DEFAULT_NUM_IMAGES)88 with gr.TabItem('Fake Scribble'):89 create_demo_fake_scribble(model.process_fake_scribble,90 max_images=MAX_IMAGES,91 default_num_images=DEFAULT_NUM_IMAGES)92 with gr.TabItem('Pose'):93 create_demo_pose(model.process_pose,94 max_images=MAX_IMAGES,95 default_num_images=DEFAULT_NUM_IMAGES)96 with gr.TabItem('Segmentation'):97 create_demo_seg(model.process_seg,98 max_images=MAX_IMAGES,99 default_num_images=DEFAULT_NUM_IMAGES)100 with gr.TabItem('Depth'):101 create_demo_depth(model.process_depth,102 max_images=MAX_IMAGES,103 default_num_images=DEFAULT_NUM_IMAGES)104 with gr.TabItem('Normal map'):105 create_demo_normal(model.process_normal,106 max_images=MAX_IMAGES,107 default_num_images=DEFAULT_NUM_IMAGES)108 109 with gr.Accordion(label='Base model', open=False):110 with gr.Row():111 with gr.Column():112 current_base_model = gr.Text(label='Current base model')113 with gr.Column(scale=0.3):114 check_base_model_button = gr.Button('Check current base model')115 with gr.Row():116 with gr.Column():117 new_base_model_id = gr.Text(118 label='New base model',119 max_lines=1,120 placeholder='runwayml/stable-diffusion-v1-5',121 info=122 'The base model must be compatible with Stable Diffusion v1.5.',123 interactive=ALLOW_CHANGING_BASE_MODEL)124 with gr.Column(scale=0.3):125 change_base_model_button = gr.Button('Change base model')126 if not ALLOW_CHANGING_BASE_MODEL:127 gr.Markdown(128 '''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>'''129 )130 131 gr.Markdown(132 '[Space using Anything-v4.0 as base model](https://huggingface.co/spaces/hysts/ControlNet-with-other-models)'133 )134 135 check_base_model_button.click(fn=lambda: model.base_model_id,136 outputs=current_base_model,137 queue=False)138 new_base_model_id.submit(fn=model.set_base_model,139 inputs=new_base_model_id,140 outputs=current_base_model)141 change_base_model_button.click(fn=model.set_base_model,142 inputs=new_base_model_id,143 outputs=current_base_model)144 145demo.queue(api_open=False).launch()146