Adapter/T2I-Adapter
169
1# demo inspired by https://huggingface.co/spaces/lambdalabs/image-mixer-demo2import argparse3import copy4import os5import shlex6import subprocess7from functools import partial8from itertools import chain9 10import cv211import gradio as gr12import torch13from basicsr.utils import tensor2img14from huggingface_hub import hf_hub_url15from pytorch_lightning import seed_everything16from torch import autocast17 18from ldm.inference_base import (DEFAULT_NEGATIVE_PROMPT, diffusion_inference, get_adapters, get_sd_models)19from ldm.modules.extra_condition import api20from ldm.modules.extra_condition.api import (ExtraCondition, get_adapter_feature, get_cond_model)21import numpy as np22from ldm.util import read_state_dict23 24torch.set_grad_enabled(False)25 26supported_cond_map = ['style', 'color', 'sketch', 'openpose', 'depth', 'canny']27supported_cond = ['style', 'color', 'sketch', 'sketch', 'openpose', 'depth', 'canny']28draw_map = gr.Interface(lambda x: x, gr.Image(source="canvas"), gr.Image())29 30# download the checkpoints31urls = {32 'TencentARC/T2I-Adapter': [33 'models/t2iadapter_keypose_sd14v1.pth', 'models/t2iadapter_color_sd14v1.pth',34 'models/t2iadapter_openpose_sd14v1.pth', 'models/t2iadapter_seg_sd14v1.pth',35 'models/t2iadapter_sketch_sd14v1.pth', 'models/t2iadapter_depth_sd14v1.pth',36 'third-party-models/body_pose_model.pth', "models/t2iadapter_style_sd14v1.pth",37 "models/t2iadapter_canny_sd14v1.pth", 'third-party-models/table5_pidinet.pth',38 "models/t2iadapter_canny_sd15v2.pth", "models/t2iadapter_depth_sd15v2.pth",39 "models/t2iadapter_sketch_sd15v2.pth"40 ],41 'runwayml/stable-diffusion-v1-5': ['v1-5-pruned-emaonly.ckpt'],42 'CompVis/stable-diffusion-v-1-4-original':['sd-v1-4.ckpt'],43 'andite/anything-v4.0': ['anything-v4.0-pruned.ckpt', 'anything-v4.0.vae.pt'],44}45 46# download image samples47torch.hub.download_url_to_file(48 'https://user-images.githubusercontent.com/52127135/223114920-cae3e723-3683-424a-bebc-0875479f2409.jpg',49 'cyber_style.jpg')50torch.hub.download_url_to_file(51 'https://user-images.githubusercontent.com/52127135/223114946-6ccc127f-cb58-443e-8677-805f5dbaf6f1.png',52 'sword.png')53torch.hub.download_url_to_file(54 'https://user-images.githubusercontent.com/52127135/223121793-20c2ac6a-5a4f-4ff8-88ea-6d007a7959dd.png',55 'white.png')56torch.hub.download_url_to_file(57 'https://user-images.githubusercontent.com/52127135/223127404-4a3748cf-85a6-40f3-af31-a74e206db96e.jpeg',58 'scream_style.jpeg')59torch.hub.download_url_to_file(60 'https://user-images.githubusercontent.com/52127135/223127433-8768913f-9872-4d24-b883-a19a3eb20623.jpg',61 'motorcycle.jpg')62 63if os.path.exists('models') == False:64 os.mkdir('models')65for repo in urls:66 files = urls[repo]67 for file in files:68 url = hf_hub_url(repo, file)69 name_ckp = url.split('/')[-1]70 save_path = os.path.join('models', name_ckp)71 if os.path.exists(save_path) == False:72 subprocess.run(shlex.split(f'wget {url} -O {save_path}'))73 74# config75parser = argparse.ArgumentParser()76parser.add_argument(77 '--sd_ckpt',78 type=str,79 default='models/v1-5-pruned-emaonly.ckpt',80 help='path to checkpoint of stable diffusion model, both .ckpt and .safetensor are supported',81)82parser.add_argument(83 '--vae_ckpt',84 type=str,85 default=None,86 help='vae checkpoint, anime SD models usually have seperate vae ckpt that need to be loaded',87)88global_opt = parser.parse_args()89global_opt.config = 'configs/stable-diffusion/sd-v1-inference.yaml'90for cond_name in supported_cond:91 if cond_name in ['sketch', 'depth', 'canny']:92 setattr(global_opt, f'{cond_name}_adapter_ckpt', f'models/t2iadapter_{cond_name}_sd15v2.pth')93 else:94 setattr(global_opt, f'{cond_name}_adapter_ckpt', f'models/t2iadapter_{cond_name}_sd14v1.pth')95global_opt.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")96global_opt.max_resolution = 512 * 51297global_opt.sampler = 'ddim'98global_opt.cond_weight = 1.099global_opt.C = 4100global_opt.f = 8101# adapters and models to processing condition inputs102adapters = {}103cond_models = {}104torch.cuda.empty_cache()105 106 107def draw_transfer(im1):108 c = im1[:, :, 0:3].astype(np.float32)109 a = im1[:, :, 3:4].astype(np.float32) / 255.0110 im1 = c * a + 255.0 * (1.0 - a)111 im1 = (im1.clip(0, 255)).astype(np.uint8)112 113 return im1114 115class process:116 def __init__(self):117 self.base_model = 'v1-5-pruned-emaonly.ckpt'118 # stable-diffusion model119 self.sd_model, self.sampler = get_sd_models(global_opt)120 121 def run(self, *args):122 opt = copy.deepcopy(global_opt)123 opt.prompt, opt.neg_prompt, opt.scale, opt.n_samples, opt.seed, opt.steps, opt.resize_short_edge, opt.cond_tau, opt.base_model \124 = args[-9:]125 # check base model126 if opt.base_model!=self.base_model:127 ckpt = os.path.join("models", opt.base_model)128 pl_sd = read_state_dict(ckpt)129 if "state_dict" in pl_sd:130 pl_sd = pl_sd["state_dict"]131 else:132 pl_sd = pl_sd133 self.sd_model.load_state_dict(pl_sd, strict=False)134 del pl_sd135 self.base_model = opt.base_model136 if self.base_model!='v1-5-pruned-emaonly.ckpt' and self.base_model!='sd-v1-4.ckpt':137 vae_sd = torch.load(os.path.join('models', 'anything-v4.0.vae.pt'), map_location="cuda")138 st = vae_sd["state_dict"]139 self.sd_model.first_stage_model.load_state_dict(st, strict=False)140 del st141 142 with torch.inference_mode(), \143 self.sd_model.ema_scope(), \144 autocast('cuda'):145 146 inps = []147 for i in range(0, len(args) - 9, len(supported_cond)):148 inps.append(args[i:i + len(supported_cond)])149 150 conds = []151 activated_conds = []152 153 ims1 = []154 ims2 = []155 for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):156 if b != 'Nothing' and (im1 is not None or im2 is not None):157 if im1 is not None and isinstance(im1,dict):158 im1 = im1['mask']159 im1 = draw_transfer(im1)160 161 if im1 is not None:162 h, w, _ = im1.shape163 else:164 h, w, _ = im2.shape165 166 # resize all the images to the same size167 for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):168 if idx == 0:169 ims1.append(im1)170 ims2.append(im2)171 continue172 if b != 'Nothing':173 if im1 is not None and isinstance(im1,dict):174 im1 = im1['mask']175 im1 = draw_transfer(im1)176 im2 = im1177 cv2.imwrite('sketch.png', im1)178 if im1 is not None:179 im1 = cv2.resize(im1, (w, h), interpolation=cv2.INTER_CUBIC)180 if im2 is not None:181 im2 = cv2.resize(im2, (w, h), interpolation=cv2.INTER_CUBIC)182 ims1.append(im1)183 ims2.append(im2)184 185 for idx, (b, _, _, cond_weight) in enumerate(zip(*inps)):186 cond_name = supported_cond[idx]187 if b == 'Nothing':188 if cond_name in adapters:189 adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)#.cpu()190 else:191 # print(idx,b)192 activated_conds.append(cond_name)193 if cond_name in adapters:194 adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)195 else:196 adapters[cond_name] = get_adapters(opt, getattr(ExtraCondition, cond_name))197 adapters[cond_name]['cond_weight'] = cond_weight198 199 process_cond_module = getattr(api, f'get_cond_{cond_name}')200 201 if b == 'Image':202 if cond_name not in cond_models:203 cond_models[cond_name] = get_cond_model(opt, getattr(ExtraCondition, cond_name))204 conds.append(process_cond_module(opt, ims1[idx], 'image', cond_models[cond_name]))205 else:206 if idx == 2: # draw207 conds.append(process_cond_module(opt, (255.-ims2[idx]).astype(np.uint8), cond_name, None))208 else:209 conds.append(process_cond_module(opt, ims2[idx], cond_name, None))210 211 adapter_features, append_to_context = get_adapter_feature(212 conds, [adapters[cond_name] for cond_name in activated_conds])213 214 output_conds = []215 for cond in conds:216 output_conds.append(tensor2img(cond, rgb2bgr=False))217 218 ims = []219 seed_everything(opt.seed)220 for _ in range(opt.n_samples):221 result = diffusion_inference(opt, self.sd_model, self.sampler, adapter_features, append_to_context)222 ims.append(tensor2img(result, rgb2bgr=False))223 224 # Clear GPU memory cache so less likely to OOM225 torch.cuda.empty_cache()226 return ims, output_conds227 228 229def change_visible(im1, im2, val):230 outputs = {}231 if val == "Image":232 outputs[im1] = gr.update(visible=True)233 outputs[im2] = gr.update(visible=False)234 elif val == "Nothing":235 outputs[im1] = gr.update(visible=False)236 outputs[im2] = gr.update(visible=False)237 else:238 outputs[im1] = gr.update(visible=False)239 outputs[im2] = gr.update(visible=True)240 return outputs241 242DESCRIPTION = '# [T2I-Adapter](https://github.com/TencentARC/T2I-Adapter)'243 244DESCRIPTION += f'<p>Gradio demo for **T2I-Adapter**: [[GitHub]](https://github.com/TencentARC/T2I-Adapter), [[Paper]](https://arxiv.org/abs/2302.08453). If T2I-Adapter is helpful, please help to ⭐ the [Github Repo](https://github.com/TencentARC/T2I-Adapter) and recommend it to your friends 😊 </p>'245 246DESCRIPTION += 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/Adapter/T2I-Adapter?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'247 248processer = process()249 250with gr.Blocks(css='style.css') as demo:251 gr.Markdown(DESCRIPTION)252 253 btns = []254 ims1 = []255 ims2 = []256 cond_weights = []257 258 with gr.Row():259 with gr.Column(scale=1.9):260 with gr.Box():261 gr.Markdown("<h5><center>Style & Color</center></h5>")262 with gr.Row():263 for cond_name in supported_cond_map[:2]:264 with gr.Box():265 with gr.Column():266 if cond_name == 'style':267 btn1 = gr.Radio(268 choices=["Image", "Nothing"],269 label=f"Input type for {cond_name}",270 interactive=True,271 value="Nothing",272 )273 else:274 btn1 = gr.Radio(275 choices=["Image", cond_name, "Nothing"],276 label=f"Input type for {cond_name}",277 interactive=True,278 value="Nothing",279 )280 281 im1 = gr.Image(282 source='upload', label="Image", interactive=True, visible=False, type="numpy")283 im2 = gr.Image(284 source='upload', label=cond_name, interactive=True, visible=False, type="numpy")285 cond_weight = gr.Slider(286 label="Condition weight",287 minimum=0,288 maximum=5,289 step=0.05,290 value=1,291 interactive=True)292 293 fn = partial(change_visible, im1, im2)294 btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)295 296 btns.append(btn1)297 ims1.append(im1)298 ims2.append(im2)299 cond_weights.append(cond_weight)300 301 with gr.Box():302 gr.Markdown("<h5><center>Drawing</center></h5>")303 with gr.Column():304 btn1 = gr.Radio(305 choices=["Sketch", "Nothing"],306 label=f"Input type for drawing",307 interactive=True,308 value="Nothing")309 im1 = gr.Image(source='canvas', tool='color-sketch', label='Pay attention to adjusting stylus thickness!', visible=False)310 im2 = im1311 cond_weight = gr.Slider(312 label="Condition weight",313 minimum=0,314 maximum=5,315 step=0.05,316 value=1,317 interactive=True)318 319 fn = partial(change_visible, im1, im2)320 btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)321 322 btns.append(btn1)323 ims1.append(im1)324 ims2.append(im2)325 cond_weights.append(cond_weight)326 327 with gr.Column(scale=4):328 with gr.Box():329 gr.Markdown("<h5><center>Structure</center></h5>")330 with gr.Row():331 for cond_name in supported_cond_map[2:6]:332 with gr.Box():333 with gr.Column():334 if cond_name == 'openpose':335 btn1 = gr.Radio(336 choices=["Image", 'pose', "Nothing"],337 label=f"Input type for {cond_name}",338 interactive=True,339 value="Nothing",340 )341 else:342 btn1 = gr.Radio(343 choices=["Image", cond_name, "Nothing"],344 label=f"Input type for {cond_name}",345 interactive=True,346 value="Nothing",347 )348 349 im1 = gr.Image(350 source='upload', label="Image", interactive=True, visible=False, type="numpy")351 im2 = gr.Image(352 source='upload', label=cond_name, interactive=True, visible=False, type="numpy")353 cond_weight = gr.Slider(354 label="Condition weight",355 minimum=0,356 maximum=5,357 step=0.05,358 value=1,359 interactive=True)360 361 fn = partial(change_visible, im1, im2)362 btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)363 btns.append(btn1)364 ims1.append(im1)365 ims2.append(im2)366 cond_weights.append(cond_weight)367 368 with gr.Column():369 base_model = gr.inputs.Radio(['v1-5-pruned-emaonly.ckpt', 'sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='v1-5-pruned-emaonly.ckpt', label='The base model you want to use. You can try more base models on https://civitai.com/.')370 prompt = gr.Textbox(label="Prompt")371 with gr.Accordion('Advanced options', open=False):372 neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)373 scale = gr.Slider(374 label="Guidance Scale (Classifier free guidance)", value=7.5, minimum=1, maximum=20, step=0.1)375 n_samples = gr.Slider(label="Num samples", value=1, minimum=1, maximum=1, step=1)376 seed = gr.Slider(label="Seed", value=42, minimum=0, maximum=10000, step=1, randomize=True)377 steps = gr.Slider(label="Steps", value=50, minimum=10, maximum=100, step=1)378 resize_short_edge = gr.Slider(label="Image resolution", value=512, minimum=320, maximum=1024, step=1)379 cond_tau = gr.Slider(380 label="timestamp parameter that determines until which step the adapter is applied",381 value=1.0,382 minimum=0.1,383 maximum=1.0,384 step=0.05)385 submit = gr.Button("Generate")386 387 with gr.Box():388 gr.Markdown("<h5><center>Results</center></h5>")389 with gr.Column():390 output = gr.Gallery().style(grid=2, height='auto')391 cond = gr.Gallery().style(grid=2, height='auto')392 393 inps = list(chain(btns, ims1, ims2, cond_weights))394 395 inps.extend([prompt, neg_prompt, scale, n_samples, seed, steps, resize_short_edge, cond_tau, base_model])396 submit.click(fn=processer.run, inputs=inps, outputs=[output, cond])397 398 ex = gr.Examples([399 [400 "Image",401 "Nothing",402 "Nothing",403 "Image",404 "Nothing",405 "Nothing",406 "Nothing",407 "cyber_style.jpg",408 "white.png",409 "white.png",410 "sword.png",411 "white.png",412 "white.png",413 "white.png",414 "white.png",415 "white.png",416 "white.png",417 "white.png",418 "white.png",419 "white.png",420 "white.png",421 1,422 1,423 1,424 1,425 1,426 1,427 1,428 "master sword",429 "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",430 7.5,431 1,432 2500,433 50,434 512,435 1,436 "v1-5-pruned-emaonly.ckpt",437 ],438 [439 "Image",440 "Nothing",441 "Nothing",442 "Image",443 "Nothing",444 "Nothing",445 "Nothing",446 "scream_style.jpeg",447 "white.png",448 "white.png",449 "motorcycle.jpg",450 "white.png",451 "white.png",452 "white.png",453 "white.png",454 "white.png",455 "white.png",456 "white.png",457 "white.png",458 "white.png",459 "white.png",460 1,461 1,462 1,463 1,464 1,465 1,466 1,467 "motorcycle",468 "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",469 7.5,470 1,471 2500,472 50,473 512,474 1,475 "v1-5-pruned-emaonly.ckpt",476 ],477 ],478 fn=processer.run,479 inputs=inps,480 outputs=[output, cond],481 cache_examples=True)482 483demo.queue().launch(debug=True, server_name='0.0.0.0')484