bedead/CoAdapter

1# demo inspired by https://huggingface.co/spaces/lambdalabs/image-mixer-demo2import argparse3import copy4import gradio as gr5import torch6from functools import partial7from itertools import chain8from torch import autocast9from pytorch_lightning import seed_everything10 11from basicsr.utils import tensor2img12from ldm.inference_base import DEFAULT_NEGATIVE_PROMPT, diffusion_inference, get_adapters, get_sd_models13from ldm.modules.extra_condition import api14from ldm.modules.extra_condition.api import ExtraCondition, get_cond_model15from ldm.modules.encoders.adapter import CoAdapterFuser16import os17from huggingface_hub import hf_hub_url18import subprocess19import shlex20import cv221 22torch.set_grad_enabled(False)23 24urls = {25    'TencentARC/T2I-Adapter':[26        'third-party-models/body_pose_model.pth', 'third-party-models/table5_pidinet.pth',27        'models/coadapter-canny-sd15v1.pth',28        'models/coadapter-color-sd15v1.pth',29        'models/coadapter-sketch-sd15v1.pth',30        'models/coadapter-style-sd15v1.pth',31        'models/coadapter-depth-sd15v1.pth',32        'models/coadapter-fuser-sd15v1.pth',33 34    ],35    'runwayml/stable-diffusion-v1-5': ['v1-5-pruned-emaonly.ckpt'],36    'andite/anything-v4.0': ['anything-v4.5-pruned.ckpt', 'anything-v4.0.vae.pt'],37}38 39if os.path.exists('models') == False:40    os.mkdir('models')41for repo in urls:42    files = urls[repo]43    for file in files:44        url = hf_hub_url(repo, file)45        name_ckp = url.split('/')[-1]46        save_path = os.path.join('models',name_ckp)47        if os.path.exists(save_path) == False:48            subprocess.run(shlex.split(f'wget {url} -O {save_path}'))49 50supported_cond = ['style', 'color', 'sketch', 'depth', 'canny']51 52# config53parser = argparse.ArgumentParser()54parser.add_argument(55    '--sd_ckpt',56    type=str,57    default='models/v1-5-pruned-emaonly.ckpt',58    help='path to checkpoint of stable diffusion model, both .ckpt and .safetensor are supported',59)60parser.add_argument(61    '--vae_ckpt',62    type=str,63    default=None,64    help='vae checkpoint, anime SD models usually have seperate vae ckpt that need to be loaded',65)66global_opt = parser.parse_args()67global_opt.config = 'configs/stable-diffusion/sd-v1-inference.yaml'68for cond_name in supported_cond:69    setattr(global_opt, f'{cond_name}_adapter_ckpt', f'models/coadapter-{cond_name}-sd15v1.pth')70global_opt.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")71global_opt.max_resolution = 512 * 51272global_opt.sampler = 'ddim'73global_opt.cond_weight = 1.074global_opt.C = 475global_opt.f = 876#TODO: expose style_cond_tau to users77global_opt.style_cond_tau = 1.078 79# stable-diffusion model80sd_model, sampler = get_sd_models(global_opt)81# adapters and models to processing condition inputs82adapters = {}83cond_models = {}84 85torch.cuda.empty_cache()86 87# fuser is indispensable88coadapter_fuser = CoAdapterFuser(unet_channels=[320, 640, 1280, 1280], width=768, num_head=8, n_layes=3)89coadapter_fuser.load_state_dict(torch.load(f'models/coadapter-fuser-sd15v1.pth'))90coadapter_fuser = coadapter_fuser.to(global_opt.device)91 92 93def run(*args):94    with torch.inference_mode(), \95            sd_model.ema_scope(), \96            autocast('cuda'):97 98        inps = []99        for i in range(0, len(args) - 8, len(supported_cond)):100            inps.append(args[i:i + len(supported_cond)])101 102        opt = copy.deepcopy(global_opt)103        opt.prompt, opt.neg_prompt, opt.scale, opt.n_samples, opt.seed, opt.steps, opt.resize_short_edge, opt.cond_tau \104            = args[-8:]105 106        ims1 = []107        ims2 = []108        for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):109            if idx > 0:110                if b != 'Nothing' and (im1 is not None or im2 is not None):111                    if im1 is not None:112                        h, w, _ = im1.shape113                    else:114                        h, w, _ = im2.shape115                # break116        # resize all the images to the same size117        for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):118            if idx == 0:119                ims1.append(im1)120                ims2.append(im2)121                continue122            if b != 'Nothing':123                if im1 is not None:124                    im1 = cv2.resize(im1, (w, h), interpolation=cv2.INTER_CUBIC)125                if im2 is not None:126                    im2 = cv2.resize(im2, (w, h), interpolation=cv2.INTER_CUBIC)127            ims1.append(im1)128            ims2.append(im2)129 130        conds = []131        activated_conds = []132        for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):133            cond_name = supported_cond[idx]134            if b == 'Nothing':135                if cond_name in adapters:136                    adapters[cond_name]['model'] = adapters[cond_name]['model'].cpu()137            else:138                activated_conds.append(cond_name)139                if cond_name in adapters:140                    adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)141                else:142                    adapters[cond_name] = get_adapters(opt, getattr(ExtraCondition, cond_name))143                adapters[cond_name]['cond_weight'] = cond_weight144 145                process_cond_module = getattr(api, f'get_cond_{cond_name}')146 147                if b == 'Image':148                    if cond_name not in cond_models:149                        cond_models[cond_name] = get_cond_model(opt, getattr(ExtraCondition, cond_name))150                    conds.append(process_cond_module(opt, ims1[idx], 'image', cond_models[cond_name]))151                else:152                    conds.append(process_cond_module(opt, ims2[idx], cond_name, None))153 154        features = dict()155        for idx, cond_name in enumerate(activated_conds):156            cur_feats = adapters[cond_name]['model'](conds[idx])157            if isinstance(cur_feats, list):158                for i in range(len(cur_feats)):159                    cur_feats[i] *= adapters[cond_name]['cond_weight']160            else:161                cur_feats *= adapters[cond_name]['cond_weight']162            features[cond_name] = cur_feats163 164        adapter_features, append_to_context = coadapter_fuser(features)165 166        output_conds = []167        for cond in conds:168            output_conds.append(tensor2img(cond, rgb2bgr=False))169 170        ims = []171        seed_everything(opt.seed)172        for _ in range(opt.n_samples):173            result = diffusion_inference(opt, sd_model, sampler, adapter_features, append_to_context)174            ims.append(tensor2img(result, rgb2bgr=False))175 176        # Clear GPU memory cache so less likely to OOM177        torch.cuda.empty_cache()178        return ims179 180 181def change_visible(im1, im2, val):182    outputs = {}183    if val == "Image":184        outputs[im1] = gr.update(visible=True)185        outputs[im2] = gr.update(visible=False)186    elif val == "Nothing":187        outputs[im1] = gr.update(visible=False)188        outputs[im2] = gr.update(visible=False)189    else:190        outputs[im1] = gr.update(visible=False)191        outputs[im2] = gr.update(visible=True)192    return outputs193 194# with gr.Blocks(title="CoAdapter", css=".gr-box {border-color: #8136e2}") as demo:195with gr.Blocks(css='style.css') as demo:196 197    btns = []198    ims1 = []199    ims2 = []200    cond_weights = []201 202    with gr.Row():203        for cond_name in supported_cond:204            with gr.Group():205                with gr.Column():206                    if cond_name == 'style':207                        btn1 = gr.Radio(208                        choices=["Image", "Nothing"],209                        label=f"Input type for {cond_name}",210                        interactive=True,211                        value="Nothing",212                    )213                    else:214                        btn1 = gr.Radio(215                            choices=["Image", cond_name, "Nothing"],216                            label=f"Input type for {cond_name}",217                            interactive=True,218                            value="Nothing",219                        )220                    im1 = gr.Image(label="Image", interactive=True, visible=False, type="numpy")221                    im2 = gr.Image(label=cond_name, interactive=True, visible=False, type="numpy")222                    cond_weight = gr.Slider(223                        label="Condition weight", minimum=0, maximum=5, step=0.05, value=1, interactive=True)224 225                    fn = partial(change_visible, im1, im2)226                    btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)227 228                    btns.append(btn1)229                    ims1.append(im1)230                    ims2.append(im2)231                    cond_weights.append(cond_weight)232 233    with gr.Column():234        prompt = gr.Textbox(label="Prompt", visible=False)235        neg_prompt = gr.Textbox(visible=False, label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)236        scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", value=7.5, minimum=1, maximum=20, step=0.1)237        n_samples = gr.Slider(label="Num samples", value=1, minimum=1, maximum=3, step=1)238        seed = gr.Slider(label="Seed", value=42, minimum=0, maximum=10000, step=1)239        steps = gr.Slider(label="Steps", value=50, minimum=10, maximum=100, step=1)240        resize_short_edge = gr.Slider(label="Image resolution", value=512, minimum=320, maximum=1024, step=1)241        cond_tau = gr.Slider(242            label="timestamp parameter that determines until which step the adapter is applied",243            value=1.0,244            minimum=0.1,245            maximum=1.0,246            step=0.05)247 248    with gr.Row():249        submit = gr.Button("Generate")250    output = gr.Gallery(rows=2, height='auto')251    # cond = gr.Gallery(rows=2, height='auto')252 253    inps = list(chain(btns, ims1, ims2, cond_weights))254    inps.extend([prompt, neg_prompt, scale, n_samples, seed, steps, resize_short_edge, cond_tau])255    submit.click(fn=run, inputs=inps, outputs=output)256# demo.launch()257demo.launch(debug=True, share=True)