Adapter/T2I-Adapter

1import cv22import os3import torch4from pytorch_lightning import seed_everything5from torch import autocast6 7from basicsr.utils import tensor2img8from ldm.inference_base import diffusion_inference, get_adapters, get_base_argument_parser, get_sd_models9from ldm.modules.extra_condition import api10from ldm.modules.extra_condition.api import ExtraCondition, get_adapter_feature, get_cond_model11 12torch.set_grad_enabled(False)13 14 15def main():16    supported_cond = [e.name for e in ExtraCondition]17    parser = get_base_argument_parser()18    for cond_name in supported_cond:19        parser.add_argument(20            f'--{cond_name}_path',21            type=str,22            default=None,23            help=f'condition image path for {cond_name}',24        )25        parser.add_argument(26            f'--{cond_name}_inp_type',27            type=str,28            default='image',29            help=f'the type of the input condition image, can be image or {cond_name}',30            choices=['image', cond_name],31        )32        parser.add_argument(33            f'--{cond_name}_adapter_ckpt',34            type=str,35            default=None,36            help=f'path to checkpoint of the {cond_name} adapter, '37                 f'if {cond_name}_path is not None, this should not be None too',38        )39        parser.add_argument(40            f'--{cond_name}_weight',41            type=float,42            default=1.0,43            help=f'the {cond_name} adapter features are multiplied by the {cond_name}_weight and then summed up together',44        )45    opt = parser.parse_args()46 47    # process argument48    activated_conds = []49    cond_paths = []50    adapter_ckpts = []51    for cond_name in supported_cond:52        if getattr(opt, f'{cond_name}_path') is None:53            continue54        assert getattr(opt, f'{cond_name}_adapter_ckpt') is not None, f'you should specify the {cond_name}_adapter_ckpt'55        activated_conds.append(cond_name)56        cond_paths.append(getattr(opt, f'{cond_name}_path'))57        adapter_ckpts.append(getattr(opt, f'{cond_name}_adapter_ckpt'))58    assert len(activated_conds) != 0, 'you did not input any condition'59 60    if opt.outdir is None:61        opt.outdir = f'outputs/test-composable-adapters'62    os.makedirs(opt.outdir, exist_ok=True)63    if opt.resize_short_edge is None:64        print(f"you don't specify the resize_shot_edge, so the maximum resolution is set to {opt.max_resolution}")65    opt.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")66 67    # prepare models68    adapters = []69    cond_models = []70    cond_inp_types = []71    process_cond_modules = []72    for cond_name in activated_conds:73        adapters.append(get_adapters(opt, getattr(ExtraCondition, cond_name)))74        cond_inp_type = getattr(opt, f'{cond_name}_inp_type', 'image')75        if cond_inp_type == 'image':76            cond_models.append(get_cond_model(opt, getattr(ExtraCondition, cond_name)))77        else:78            cond_models.append(None)79        cond_inp_types.append(cond_inp_type)80        process_cond_modules.append(getattr(api, f'get_cond_{cond_name}'))81    sd_model, sampler = get_sd_models(opt)82 83    # inference84    with torch.inference_mode(), \85            sd_model.ema_scope(), \86            autocast('cuda'):87        seed_everything(opt.seed)88        conds = []89        for cond_idx, cond_name in enumerate(activated_conds):90            conds.append(process_cond_modules[cond_idx](91                opt, cond_paths[cond_idx], cond_inp_types[cond_idx], cond_models[cond_idx],92            ))93        adapter_features, append_to_context = get_adapter_feature(conds, adapters)94        for v_idx in range(opt.n_samples):95            result = diffusion_inference(opt, sd_model, sampler, adapter_features, append_to_context)96            base_count = len(os.listdir(opt.outdir))97            cv2.imwrite(os.path.join(opt.outdir, f'{base_count:05}_result.png'), tensor2img(result))98 99 100if __name__ == '__main__':101    main()102