upload .py files at root dir

regionally_controlable_sampling.py

@@ -0,0 +1,189 @@
+import argparse
+import hashlib
+import json
+import os.path
+
+import torch
+from diffusers import DPMSolverMultistepScheduler
+from diffusers.models import T2IAdapter
+from PIL import Image
+
+from mixofshow.pipelines.pipeline_regionally_t2iadapter import RegionallyT2IAdapterPipeline
+
+
+def sample_image(pipe,
+    input_prompt,
+    input_neg_prompt=None,
+    generator=None,
+    num_inference_steps=50,
+    guidance_scale=7.5,
+    sketch_adaptor_weight=1.0,
+    region_sketch_adaptor_weight='',
+    keypose_adaptor_weight=1.0,
+    region_keypose_adaptor_weight='',
+    **extra_kargs
+):
+
+    keypose_condition = extra_kargs.pop('keypose_condition')
+    if keypose_condition is not None:
+        keypose_adapter_input = [keypose_condition] * len(input_prompt)
+    else:
+        keypose_adapter_input = None
+
+    sketch_condition = extra_kargs.pop('sketch_condition')
+    if sketch_condition is not None:
+        sketch_adapter_input = [sketch_condition] * len(input_prompt)
+    else:
+        sketch_adapter_input = None
+
+    images = pipe(
+        prompt=input_prompt,
+        negative_prompt=input_neg_prompt,
+        keypose_adapter_input=keypose_adapter_input,
+        keypose_adaptor_weight=keypose_adaptor_weight,
+        region_keypose_adaptor_weight=region_keypose_adaptor_weight,
+        sketch_adapter_input=sketch_adapter_input,
+        sketch_adaptor_weight=sketch_adaptor_weight,
+        region_sketch_adaptor_weight=region_sketch_adaptor_weight,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        **extra_kargs).images
+    return images
+
+
+def build_model(pretrained_model, device):
+    pipe = RegionallyT2IAdapterPipeline.from_pretrained(pretrained_model, torch_dtype=torch.float16).to(device)
+    assert os.path.exists(os.path.join(pretrained_model, 'new_concept_cfg.json'))
+    with open(os.path.join(pretrained_model, 'new_concept_cfg.json'), 'r') as json_file:
+        new_concept_cfg = json.load(json_file)
+    pipe.set_new_concept_cfg(new_concept_cfg)
+    pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained(pretrained_model, subfolder='scheduler')
+    pipe.keypose_adapter = T2IAdapter.from_pretrained('TencentARC/t2iadapter_openpose_sd14v1', torch_dtype=torch.float16).to(device)
+    pipe.sketch_adapter = T2IAdapter.from_pretrained('TencentARC/t2iadapter_sketch_sd14v1', torch_dtype=torch.float16).to(device)
+    return pipe
+
+
+def prepare_text(prompt, region_prompts, height, width):
+    '''
+    Args:
+        prompt_entity: [subject1]-*-[attribute1]-*-[Location1]|[subject2]-*-[attribute2]-*-[Location2]|[global text]
+    Returns:
+        full_prompt: subject1, attribute1 and subject2, attribute2, global text
+        context_prompt: subject1 and subject2, global text
+        entity_collection: [(subject1, attribute1), Location1]
+    '''
+    region_collection = []
+
+    regions = region_prompts.split('|')
+
+    for region in regions:
+        if region == '':
+            break
+        prompt_region, neg_prompt_region, pos = region.split('-*-')
+        prompt_region = prompt_region.replace('[', '').replace(']', '')
+        neg_prompt_region = neg_prompt_region.replace('[', '').replace(']', '')
+        pos = eval(pos)
+        if len(pos) == 0:
+            pos = [0, 0, 1, 1]
+        else:
+            pos[0], pos[2] = pos[0] / height, pos[2] / height
+            pos[1], pos[3] = pos[1] / width, pos[3] / width
+
+        region_collection.append((prompt_region, neg_prompt_region, pos))
+    return (prompt, region_collection)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser('', add_help=False)
+    parser.add_argument('--pretrained_model', default='experiments/composed_edlora/anythingv4/hina+kario+tezuka+mitsuha+son_anythingv4/combined_model_base', type=str)
+    parser.add_argument('--sketch_condition', default=None, type=str)
+    parser.add_argument('--sketch_adaptor_weight', default=1.0, type=float)
+    parser.add_argument('--region_sketch_adaptor_weight', default='', type=str)
+    parser.add_argument('--keypose_condition', default=None, type=str)
+    parser.add_argument('--keypose_adaptor_weight', default=1.0, type=float)
+    parser.add_argument('--region_keypose_adaptor_weight', default='', type=str)
+    parser.add_argument('--save_dir', default=None, type=str)
+    parser.add_argument('--prompt', default='photo of a toy', type=str)
+    parser.add_argument('--negative_prompt', default='', type=str)
+    parser.add_argument('--prompt_rewrite', default='', type=str)
+    parser.add_argument('--seed', default=16141, type=int)
+    parser.add_argument('--suffix', default='', type=str)
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    pipe = build_model(args.pretrained_model, device)
+
+    if args.sketch_condition is not None and os.path.exists(args.sketch_condition):
+        sketch_condition = Image.open(args.sketch_condition).convert('L')
+        width_sketch, height_sketch = sketch_condition.size
+        print('use sketch condition')
+    else:
+        sketch_condition, width_sketch, height_sketch = None, 0, 0
+        print('skip sketch condition')
+
+    if args.keypose_condition is not None and os.path.exists(args.keypose_condition):
+        keypose_condition = Image.open(args.keypose_condition).convert('RGB')
+        width_pose, height_pose = keypose_condition.size
+        print('use pose condition')
+    else:
+        keypose_condition, width_pose, height_pose = None, 0, 0
+        print('skip pose condition')
+
+    if width_sketch != 0 and width_pose != 0:
+        assert width_sketch == width_pose and height_sketch == height_pose, 'conditions should be same size'
+    width, height = max(width_pose, width_sketch), max(height_pose, height_sketch)
+
+    kwargs = {
+        'sketch_condition': sketch_condition,
+        'keypose_condition': keypose_condition,
+        'height': height,
+        'width': width,
+    }
+
+    prompts = [args.prompt]
+    prompts_rewrite = [args.prompt_rewrite]
+    input_prompt = [prepare_text(p, p_w, height, width) for p, p_w in zip(prompts, prompts_rewrite)]
+    save_prompt = input_prompt[0][0]
+
+    image = sample_image(
+        pipe,
+        input_prompt=input_prompt,
+        input_neg_prompt=[args.negative_prompt] * len(input_prompt),
+        generator=torch.Generator(device).manual_seed(args.seed),
+        sketch_adaptor_weight=args.sketch_adaptor_weight,
+        region_sketch_adaptor_weight=args.region_sketch_adaptor_weight,
+        keypose_adaptor_weight=args.keypose_adaptor_weight,
+        region_keypose_adaptor_weight=args.region_keypose_adaptor_weight,
+        **kwargs)
+
+    print(f'save to: {args.save_dir}')
+
+    configs = [
+        f'pretrained_model: {args.pretrained_model}\n',
+        f'context_prompt: {args.prompt}\n', f'neg_context_prompt: {args.negative_prompt}\n',
+        f'sketch_condition: {args.sketch_condition}\n', f'sketch_adaptor_weight: {args.sketch_adaptor_weight}\n',
+        f'region_sketch_adaptor_weight: {args.region_sketch_adaptor_weight}\n',
+        f'keypose_condition: {args.keypose_condition}\n', f'keypose_adaptor_weight: {args.keypose_adaptor_weight}\n',
+        f'region_keypose_adaptor_weight: {args.region_keypose_adaptor_weight}\n', f'random seed: {args.seed}\n',
+        f'prompt_rewrite: {args.prompt_rewrite}\n'
+    ]
+    hash_code = hashlib.sha256(''.join(configs).encode('utf-8')).hexdigest()[:8]
+
+    save_prompt = save_prompt.replace(' ', '_')
+    # save_name = f'{save_prompt}---{args.suffix}---{hash_code}.png'
+    # save_dir = os.path.join(args.save_dir, f'seed_{args.seed}')
+    save_name = f'{save_prompt}---{args.suffix}(seed{args.seed})---{hash_code}.png'
+    save_dir = args.save_dir
+    save_path = os.path.join(save_dir, save_name)
+    save_config_path = os.path.join(save_dir, save_name.replace('.png', '.txt'))
+
+    os.makedirs(save_dir, exist_ok=True)
+    image[0].save(os.path.join(save_dir, save_name))
+
+    with open(save_config_path, 'w') as fw:
+        fw.writelines(configs)

test_edlora.py

@@ -0,0 +1,110 @@
+import argparse
+import os
+import os.path as osp
+
+import torch
+import torch.utils.checkpoint
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from diffusers import DPMSolverMultistepScheduler
+from diffusers.utils import check_min_version
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from mixofshow.data.prompt_dataset import PromptDataset
+from mixofshow.pipelines.pipeline_edlora import EDLoRAPipeline, StableDiffusionPipeline
+from mixofshow.utils.convert_edlora_to_diffusers import convert_edlora
+from mixofshow.utils.util import NEGATIVE_PROMPT, compose_visualize, dict2str, pil_imwrite, set_path_logger
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version('0.18.2')
+
+
+def visual_validation(accelerator, pipe, dataloader, current_iter, opt):
+    dataset_name = dataloader.dataset.opt['name']
+    pipe.unet.eval()
+    pipe.text_encoder.eval()
+
+    for idx, val_data in enumerate(tqdm(dataloader)):
+        output = pipe(
+            prompt=val_data['prompts'],
+            latents=val_data['latents'].to(dtype=torch.float16),
+            negative_prompt=[NEGATIVE_PROMPT] * len(val_data['prompts']),
+            num_inference_steps=opt['val']['sample'].get('num_inference_steps', 50),
+            guidance_scale=opt['val']['sample'].get('guidance_scale', 7.5),
+        ).images
+
+        for img, prompt, indice in zip(output, val_data['prompts'], val_data['indices']):
+            img_name = '{prompt}---G_{guidance_scale}_S_{steps}---{indice}'.format(
+                prompt=prompt.replace(' ', '_'),
+                guidance_scale=opt['val']['sample'].get('guidance_scale', 7.5),
+                steps=opt['val']['sample'].get('num_inference_steps', 50),
+                indice=indice)
+
+            save_img_path = osp.join(opt['path']['visualization'], dataset_name, f'{current_iter}', f'{img_name}---{current_iter}.png')
+
+            pil_imwrite(img, save_img_path)
+        # tentative for out of GPU memory
+        del output
+        torch.cuda.empty_cache()
+
+    # Save the lora layers, final eval
+    accelerator.wait_for_everyone()
+
+    if opt['val'].get('compose_visualize'):
+        if accelerator.is_main_process:
+            compose_visualize(os.path.dirname(save_img_path))
+
+
+def test(root_path, args):
+
+    # load config
+    opt = OmegaConf.to_container(OmegaConf.load(args.opt), resolve=True)
+
+    # set accelerator, mix-precision set in the environment by "accelerate config"
+    accelerator = Accelerator(mixed_precision=opt['mixed_precision'])
+
+    # set experiment dir
+    with accelerator.main_process_first():
+        set_path_logger(accelerator, root_path, args.opt, opt, is_train=False)
+
+    # get logger
+    logger = get_logger('mixofshow', log_level='INFO')
+    logger.info(accelerator.state, main_process_only=True)
+
+    logger.info(dict2str(opt))
+
+    # If passed along, set the training seed now.
+    if opt.get('manual_seed') is not None:
+        set_seed(opt['manual_seed'])
+
+    # Get the training dataset
+    valset_cfg = opt['datasets']['val_vis']
+    val_dataset = PromptDataset(valset_cfg)
+    val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=valset_cfg['batch_size_per_gpu'], shuffle=False)
+
+    enable_edlora = opt['models']['enable_edlora']
+
+    for lora_alpha in opt['val']['alpha_list']:
+        pipeclass = EDLoRAPipeline if enable_edlora else StableDiffusionPipeline
+        pipe = pipeclass.from_pretrained(opt['models']['pretrained_path'],
+            scheduler=DPMSolverMultistepScheduler.from_pretrained(opt['models']['pretrained_path'], subfolder='scheduler'),
+            torch_dtype=torch.float16).to('cuda')
+        pipe, new_concept_cfg = convert_edlora(pipe, torch.load(opt['path']['lora_path']), enable_edlora=enable_edlora, alpha=lora_alpha)
+        pipe.set_new_concept_cfg(new_concept_cfg)
+        # visualize embedding + LoRA weight shift
+        logger.info(f'Start validation sample lora({lora_alpha}):')
+
+        lora_type = 'edlora' if enable_edlora else 'lora'
+        visual_validation(accelerator, pipe, val_dataloader, f'validation_{lora_type}_{lora_alpha}', opt)
+        del pipe
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, default='options/test/EDLoRA/EDLoRA_hina_Anyv4_B4_Iter1K.yml')
+    args = parser.parse_args()
+
+    root_path = osp.abspath(osp.join(__file__, osp.pardir))
+    test(root_path, args)

train_edlora.py

@@ -0,0 +1,198 @@
+import argparse
+import copy
+import os
+import os.path as osp
+
+import torch
+import torch.utils.checkpoint
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from diffusers import DPMSolverMultistepScheduler
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version
+from omegaconf import OmegaConf
+
+from mixofshow.data.lora_dataset import LoraDataset
+from mixofshow.data.prompt_dataset import PromptDataset
+from mixofshow.pipelines.pipeline_edlora import EDLoRAPipeline, StableDiffusionPipeline
+from mixofshow.pipelines.trainer_edlora import EDLoRATrainer
+from mixofshow.utils.convert_edlora_to_diffusers import convert_edlora
+from mixofshow.utils.util import MessageLogger, dict2str, reduce_loss_dict, set_path_logger
+from test_edlora import visual_validation
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version('0.18.2')
+
+
+def train(root_path, args):
+
+    # load config
+    opt = OmegaConf.to_container(OmegaConf.load(args.opt), resolve=True)
+
+    # set accelerator, mix-precision set in the environment by "accelerate config"
+    accelerator = Accelerator(mixed_precision=opt['mixed_precision'], gradient_accumulation_steps=opt['gradient_accumulation_steps'])
+
+    # set experiment dir
+    with accelerator.main_process_first():
+        set_path_logger(accelerator, root_path, args.opt, opt, is_train=True)
+
+    # get logger
+    logger = get_logger('mixofshow', log_level='INFO')
+    logger.info(accelerator.state, main_process_only=True)
+
+    logger.info(dict2str(opt))
+
+    # If passed along, set the training seed now.
+    if opt.get('manual_seed') is not None:
+        set_seed(opt['manual_seed'])
+
+    # Load model
+    EDLoRA_trainer = EDLoRATrainer(**opt['models'])
+
+    # set optimizer
+    train_opt = opt['train']
+    optim_type = train_opt['optim_g'].pop('type')
+    assert optim_type == 'AdamW', 'only support AdamW now'
+    optimizer = torch.optim.AdamW(EDLoRA_trainer.get_params_to_optimize(), **train_opt['optim_g'])
+
+    # Get the training dataset
+    trainset_cfg = opt['datasets']['train']
+    train_dataset = LoraDataset(trainset_cfg)
+    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=trainset_cfg['batch_size_per_gpu'], shuffle=True, drop_last=True)
+
+    # Get the training dataset
+    valset_cfg = opt['datasets']['val_vis']
+    val_dataset = PromptDataset(valset_cfg)
+    val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=valset_cfg['batch_size_per_gpu'], shuffle=False)
+
+    # Prepare everything with our `accelerator`.
+    EDLoRA_trainer, optimizer, train_dataloader, val_dataloader = accelerator.prepare(EDLoRA_trainer, optimizer, train_dataloader, val_dataloader)
+
+    # Train!
+    total_batch_size = opt['datasets']['train']['batch_size_per_gpu'] * accelerator.num_processes * opt['gradient_accumulation_steps']
+    total_iter = len(train_dataset) / total_batch_size
+    opt['train']['total_iter'] = total_iter
+
+    logger.info('***** Running training *****')
+    logger.info(f'  Num examples = {len(train_dataset)}')
+    logger.info(f"  Instantaneous batch size per device = {opt['datasets']['train']['batch_size_per_gpu']}")
+    logger.info(f'  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}')
+    logger.info(f'  Total optimization steps = {total_iter}')
+    global_step = 0
+
+    # Scheduler
+    lr_scheduler = get_scheduler(
+        'linear',
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        num_training_steps=total_iter * opt['gradient_accumulation_steps'],
+    )
+
+    def make_data_yielder(dataloader):
+        while True:
+            for batch in dataloader:
+                yield batch
+            accelerator.wait_for_everyone()
+
+    train_data_yielder = make_data_yielder(train_dataloader)
+
+    msg_logger = MessageLogger(opt, global_step)
+    stop_emb_update = False
+
+    original_embedding = copy.deepcopy(accelerator.unwrap_model(EDLoRA_trainer).text_encoder.get_input_embeddings().weight)
+
+    while global_step < opt['train']['total_iter']:
+        with accelerator.accumulate(EDLoRA_trainer):
+
+            accelerator.unwrap_model(EDLoRA_trainer).unet.train()
+            accelerator.unwrap_model(EDLoRA_trainer).text_encoder.train()
+            loss_dict = {}
+
+            batch = next(train_data_yielder)
+
+            if 'masks' in batch:
+                masks = batch['masks']
+            else:
+                masks = batch['img_masks']
+
+            loss = EDLoRA_trainer(batch['images'], batch['prompts'], masks, batch['img_masks'])
+            loss_dict['loss'] = loss
+
+            # get fix embedding and learn embedding
+            index_no_updates = torch.arange(len(accelerator.unwrap_model(EDLoRA_trainer).tokenizer)) != -1
+            if not stop_emb_update:
+                for token_id in accelerator.unwrap_model(EDLoRA_trainer).get_all_concept_token_ids():
+                    index_no_updates[token_id] = False
+
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        if accelerator.sync_gradients:
+            # set no update token to origin
+            token_embeds = accelerator.unwrap_model(EDLoRA_trainer).text_encoder.get_input_embeddings().weight
+            token_embeds.data[index_no_updates, :] = original_embedding.data[index_no_updates, :]
+
+            token_embeds = accelerator.unwrap_model(EDLoRA_trainer).text_encoder.get_input_embeddings().weight
+            concept_token_ids = accelerator.unwrap_model(EDLoRA_trainer).get_all_concept_token_ids()
+            loss_dict['Norm_mean'] = token_embeds[concept_token_ids].norm(dim=-1).mean()
+            if stop_emb_update is False and float(loss_dict['Norm_mean']) >= train_opt.get('emb_norm_threshold', 5.5e-1):
+                stop_emb_update = True
+                original_embedding = copy.deepcopy(accelerator.unwrap_model(EDLoRA_trainer).text_encoder.get_input_embeddings().weight)
+
+            log_dict = reduce_loss_dict(accelerator, loss_dict)
+
+        # Checks if the accelerator has performed an optimization step behind the scenes
+        if accelerator.sync_gradients:
+            global_step += 1
+
+            if global_step % opt['logger']['print_freq'] == 0:
+                log_vars = {'iter': global_step}
+                log_vars.update({'lrs': lr_scheduler.get_last_lr()})
+                log_vars.update(log_dict)
+                msg_logger(log_vars)
+
+            if global_step % opt['logger']['save_checkpoint_freq'] == 0:
+                save_and_validation(accelerator, opt, EDLoRA_trainer, val_dataloader, global_step, logger)
+
+    # Save the lora layers, final eval
+    accelerator.wait_for_everyone()
+    save_and_validation(accelerator, opt, EDLoRA_trainer, val_dataloader, 'latest', logger)
+
+
+def save_and_validation(accelerator, opt, EDLoRA_trainer, val_dataloader, global_step, logger):
+    enable_edlora = opt['models']['enable_edlora']
+    lora_type = 'edlora' if enable_edlora else 'lora'
+    save_path = os.path.join(opt['path']['models'], f'{lora_type}_model-{global_step}.pth')
+
+    if accelerator.is_main_process:
+        accelerator.save({'params': accelerator.unwrap_model(EDLoRA_trainer).delta_state_dict()}, save_path)
+        logger.info(f'Save state to {save_path}')
+
+    accelerator.wait_for_everyone()
+
+    if opt['val']['val_during_save']:
+        logger.info(f'Start validation {save_path}:')
+        for lora_alpha in opt['val']['alpha_list']:
+            pipeclass = EDLoRAPipeline if enable_edlora else StableDiffusionPipeline
+
+            pipe = pipeclass.from_pretrained(opt['models']['pretrained_path'],
+                scheduler=DPMSolverMultistepScheduler.from_pretrained(opt['models']['pretrained_path'], subfolder='scheduler'),
+                torch_dtype=torch.float16).to('cuda')
+            pipe, new_concept_cfg = convert_edlora(pipe, torch.load(save_path), enable_edlora=enable_edlora, alpha=lora_alpha)
+            pipe.set_new_concept_cfg(new_concept_cfg)
+            pipe.set_progress_bar_config(disable=True)
+            visual_validation(accelerator, pipe, val_dataloader, f'Iters-{global_step}_Alpha-{lora_alpha}', opt)
+
+            del pipe
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, default='options/train/EDLoRA/EDLoRA_hina_Anyv4_B4_Iter1K.yml')
+    args = parser.parse_args()
+
+    root_path = osp.abspath(osp.join(__file__, osp.pardir))
+    train(root_path, args)

weight_fusion.py

@@ -0,0 +1,699 @@
+import argparse
+import copy
+import itertools
+import json
+import logging
+import os
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+from diffusers import DDPMScheduler, DPMSolverMultistepScheduler, StableDiffusionPipeline
+from tqdm import tqdm
+
+from mixofshow.models.edlora import revise_edlora_unet_attention_forward
+from mixofshow.pipelines.pipeline_edlora import bind_concept_prompt
+from mixofshow.utils.util import set_logger
+
+TEMPLATE_SIMPLE = 'photo of a {}'
+
+
+def chunk_compute_mse(K_target, V_target, W, device, chunk_size=5000):
+    num_chunks = (K_target.size(0) + chunk_size - 1) // chunk_size
+
+    loss = 0
+
+    for i in range(num_chunks):
+        # Extract the current chunk
+        start_idx = i * chunk_size
+        end_idx = min(start_idx + chunk_size, K_target.size(0))
+        loss += F.mse_loss(
+            F.linear(K_target[start_idx:end_idx].to(device), W),
+            V_target[start_idx:end_idx].to(device)) * (end_idx - start_idx)
+    loss /= K_target.size(0)
+    return loss
+
+
+def update_quasi_newton(K_target, V_target, W, iters, device):
+    '''
+    Args:
+        K: torch.Tensor, size [n_samples, n_features]
+        V: torch.Tensor, size [n_samples, n_targets]
+        K_target: torch.Tensor, size [n_constraints, n_features]
+        V_target: torch.Tensor, size [n_constraints, n_targets]
+        W: torch.Tensor, size [n_targets, n_features]
+
+    Returns:
+        Wnew: torch.Tensor, size [n_targets, n_features]
+    '''
+
+    W = W.detach()
+    V_target = V_target.detach()
+    K_target = K_target.detach()
+
+    W.requires_grad = True
+    K_target.requires_grad = False
+    V_target.requires_grad = False
+
+    best_loss = np.Inf
+    best_W = None
+
+    def closure():
+        nonlocal best_W, best_loss
+        optimizer.zero_grad()
+
+        if len(W.shape) == 4:
+            loss = F.mse_loss(F.conv2d(K_target.to(device), W),
+                              V_target.to(device))
+        else:
+            loss = chunk_compute_mse(K_target, V_target, W, device)
+
+        if loss < best_loss:
+            best_loss = loss
+            best_W = W.clone().cpu()
+        loss.backward()
+        return loss
+
+    optimizer = optim.LBFGS([W],
+                            lr=1,
+                            max_iter=iters,
+                            history_size=25,
+                            line_search_fn='strong_wolfe',
+                            tolerance_grad=1e-16,
+                            tolerance_change=1e-16)
+    optimizer.step(closure)
+
+    with torch.no_grad():
+        if len(W.shape) == 4:
+            loss = torch.norm(
+                F.conv2d(K_target.to(torch.float32), best_W.to(torch.float32)) - V_target.to(torch.float32), 2, dim=1)
+        else:
+            loss = torch.norm(
+                F.linear(K_target.to(torch.float32), best_W.to(torch.float32)) - V_target.to(torch.float32), 2, dim=1)
+
+    logging.info('new_concept loss: %e' % loss.mean().item())
+    return best_W
+
+
+def merge_lora_into_weight(original_state_dict, lora_state_dict, modification_layer_names, model_type, alpha, device):
+    def get_lora_down_name(original_layer_name):
+        if model_type == 'text_encoder':
+            lora_down_name = original_layer_name.replace('q_proj.weight', 'q_proj.lora_down.weight') \
+                .replace('k_proj.weight', 'k_proj.lora_down.weight') \
+                .replace('v_proj.weight', 'v_proj.lora_down.weight') \
+                .replace('out_proj.weight', 'out_proj.lora_down.weight') \
+                .replace('fc1.weight', 'fc1.lora_down.weight') \
+                .replace('fc2.weight', 'fc2.lora_down.weight')
+        else:
+            lora_down_name = k.replace('to_q.weight', 'to_q.lora_down.weight') \
+                .replace('to_k.weight', 'to_k.lora_down.weight') \
+                .replace('to_v.weight', 'to_v.lora_down.weight') \
+                .replace('to_out.0.weight', 'to_out.0.lora_down.weight') \
+                .replace('ff.net.0.proj.weight', 'ff.net.0.proj.lora_down.weight') \
+                .replace('ff.net.2.weight', 'ff.net.2.lora_down.weight') \
+                .replace('proj_out.weight', 'proj_out.lora_down.weight') \
+                .replace('proj_in.weight', 'proj_in.lora_down.weight')
+
+        return lora_down_name
+
+    assert model_type in ['unet', 'text_encoder']
+    new_state_dict = copy.deepcopy(original_state_dict)
+    load_cnt = 0
+
+    for k in modification_layer_names:
+        lora_down_name = get_lora_down_name(k)
+        lora_up_name = lora_down_name.replace('lora_down', 'lora_up')
+
+        if lora_up_name in lora_state_dict:
+            load_cnt += 1
+            original_params = new_state_dict[k]
+            lora_down_params = lora_state_dict[lora_down_name].to(device)
+            lora_up_params = lora_state_dict[lora_up_name].to(device)
+            if len(original_params.shape) == 4:
+                lora_param = lora_up_params.squeeze(
+                ) @ lora_down_params.squeeze()
+                lora_param = lora_param.unsqueeze(-1).unsqueeze(-1)
+            else:
+                lora_param = lora_up_params @ lora_down_params
+            merge_params = original_params + alpha * lora_param
+            new_state_dict[k] = merge_params
+
+    logging.info(f'load {load_cnt} LoRAs of {model_type}')
+    return new_state_dict
+
+
+module_io_recoder = {}
+record_feature = False  # remember to set record feature
+
+
+def get_hooker(module_name):
+    def hook(module, feature_in, feature_out):
+        if module_name not in module_io_recoder:
+            module_io_recoder[module_name] = {'input': [], 'output': []}
+        if record_feature:
+            module_io_recoder[module_name]['input'].append(feature_in[0].cpu())
+            if module.bias is not None:
+                if len(feature_out.shape) == 4:
+                    bias = module.bias.unsqueeze(-1).unsqueeze(-1)
+                else:
+                    bias = module.bias
+                module_io_recoder[module_name]['output'].append(
+                    (feature_out - bias).cpu())  # remove bias
+            else:
+                module_io_recoder[module_name]['output'].append(
+                    feature_out.cpu())
+
+    return hook
+
+
+def init_stable_diffusion(pretrained_model_path, device):
+    # step1: get w0 parameters
+    model_id = pretrained_model_path
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
+
+    train_scheduler = DDPMScheduler.from_pretrained(model_id, subfolder='scheduler')
+    test_scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder='scheduler')
+    pipe.safety_checker = None
+    pipe.scheduler = test_scheduler
+    return pipe, train_scheduler, test_scheduler
+
+
+@torch.no_grad()
+def get_text_feature(prompts, tokenizer, text_encoder, device, return_type='category_embedding'):
+    text_features = []
+
+    if return_type == 'category_embedding':
+        for text in prompts:
+            tokens = tokenizer(
+                text,
+                truncation=True,
+                max_length=tokenizer.model_max_length,
+                return_length=True,
+                return_overflowing_tokens=False,
+                padding='do_not_pad',
+            ).input_ids
+
+            new_token_position = torch.where(torch.tensor(tokens) >= 49407)[0]
+            # >40497 not include end token | >=40497 include end token
+            concept_feature = text_encoder(
+                torch.LongTensor(tokens).reshape(
+                    1, -1).to(device))[0][:,
+                              new_token_position].reshape(-1, 768)
+            text_features.append(concept_feature)
+        return torch.cat(text_features, 0).float()
+    elif return_type == 'full_embedding':
+        text_input = tokenizer(prompts,
+                               padding='max_length',
+                               max_length=tokenizer.model_max_length,
+                               truncation=True,
+                               return_tensors='pt')
+        text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
+        return text_embeddings
+    else:
+        raise NotImplementedError
+
+
+def merge_new_concepts_(embedding_list, concept_list, tokenizer, text_encoder):
+    def add_new_concept(concept_name, embedding):
+        new_token_names = [
+            f'<new{start_idx + layer_id}>'
+            for layer_id in range(NUM_CROSS_ATTENTION_LAYERS)
+        ]
+        num_added_tokens = tokenizer.add_tokens(new_token_names)
+        assert num_added_tokens == NUM_CROSS_ATTENTION_LAYERS
+        new_token_ids = [
+            tokenizer.convert_tokens_to_ids(token_name)
+            for token_name in new_token_names
+        ]
+
+        text_encoder.resize_token_embeddings(len(tokenizer))
+        token_embeds = text_encoder.get_input_embeddings().weight.data
+
+        token_embeds[new_token_ids] = token_embeds[new_token_ids].copy_(
+            embedding[concept_name])
+
+        embedding_features.update({concept_name: embedding[concept_name]})
+        logging.info(
+            f'concept {concept_name} is bind with token_id: [{min(new_token_ids)}, {max(new_token_ids)}]'
+        )
+
+        return start_idx + NUM_CROSS_ATTENTION_LAYERS, new_token_ids, new_token_names
+
+    embedding_features = {}
+    new_concept_cfg = {}
+
+    start_idx = 0
+
+    NUM_CROSS_ATTENTION_LAYERS = 16
+
+    for idx, (embedding,
+              concept) in enumerate(zip(embedding_list, concept_list)):
+        concept_names = concept['concept_name'].split(' ')
+
+        for concept_name in concept_names:
+            if not concept_name.startswith('<'):
+                continue
+            else:
+                assert concept_name in embedding, 'check the config, the provide concept name is not in the lora model'
+            start_idx, new_token_ids, new_token_names = add_new_concept(
+                concept_name, embedding)
+            new_concept_cfg.update({
+                concept_name: {
+                    'concept_token_ids': new_token_ids,
+                    'concept_token_names': new_token_names
+                }
+            })
+    return embedding_features, new_concept_cfg
+
+
+def parse_new_concepts(concept_cfg):
+    with open(concept_cfg, 'r') as f:
+        concept_list = json.load(f)
+
+    model_paths = [concept['lora_path'] for concept in concept_list]
+
+    embedding_list = []
+    text_encoder_list = []
+    unet_crosskv_list = []
+    unet_spatial_attn_list = []
+
+    for model_path in model_paths:
+        model = torch.load(model_path)['params']
+
+        if 'new_concept_embedding' in model and len(
+                model['new_concept_embedding']) != 0:
+            embedding_list.append(model['new_concept_embedding'])
+        else:
+            embedding_list.append(None)
+
+        if 'text_encoder' in model and len(model['text_encoder']) != 0:
+            text_encoder_list.append(model['text_encoder'])
+        else:
+            text_encoder_list.append(None)
+
+        if 'unet' in model and len(model['unet']) != 0:
+            crosskv_matches = ['attn2.to_k.lora', 'attn2.to_v.lora']
+            crosskv_dict = {
+                k: v
+                for k, v in model['unet'].items()
+                if any([x in k for x in crosskv_matches])
+            }
+
+            if len(crosskv_dict) != 0:
+                unet_crosskv_list.append(crosskv_dict)
+            else:
+                unet_crosskv_list.append(None)
+
+            spatial_attn_dict = {
+                k: v
+                for k, v in model['unet'].items()
+                if all([x not in k for x in crosskv_matches])
+            }
+
+            if len(spatial_attn_dict) != 0:
+                unet_spatial_attn_list.append(spatial_attn_dict)
+            else:
+                unet_spatial_attn_list.append(None)
+        else:
+            unet_crosskv_list.append(None)
+            unet_spatial_attn_list.append(None)
+
+    return embedding_list, text_encoder_list, unet_crosskv_list, unet_spatial_attn_list, concept_list
+
+
+def merge_kv_in_cross_attention(concept_list, optimize_iters, new_concept_cfg,
+                                tokenizer, text_encoder, unet,
+                                unet_crosskv_list, device):
+    # crosskv attention layer names
+    matches = ['attn2.to_k', 'attn2.to_v']
+
+    cross_attention_idx = -1
+    cross_kv_layer_names = []
+
+    # the crosskv name should match the order down->mid->up, and record its layer id
+    for name, _ in unet.down_blocks.named_parameters():
+        if any([x in name for x in matches]):
+            if 'to_k' in name:
+                cross_attention_idx += 1
+                cross_kv_layer_names.append(
+                    (cross_attention_idx, 'down_blocks.' + name))
+                cross_kv_layer_names.append(
+                    (cross_attention_idx,
+                     'down_blocks.' + name.replace('to_k', 'to_v')))
+            else:
+                pass
+
+    for name, _ in unet.mid_block.named_parameters():
+        if any([x in name for x in matches]):
+            if 'to_k' in name:
+                cross_attention_idx += 1
+                cross_kv_layer_names.append(
+                    (cross_attention_idx, 'mid_block.' + name))
+                cross_kv_layer_names.append(
+                    (cross_attention_idx,
+                     'mid_block.' + name.replace('to_k', 'to_v')))
+            else:
+                pass
+
+    for name, _ in unet.up_blocks.named_parameters():
+        if any([x in name for x in matches]):
+            if 'to_k' in name:
+                cross_attention_idx += 1
+                cross_kv_layer_names.append(
+                    (cross_attention_idx, 'up_blocks.' + name))
+                cross_kv_layer_names.append(
+                    (cross_attention_idx,
+                     'up_blocks.' + name.replace('to_k', 'to_v')))
+            else:
+                pass
+
+    logging.info(
+        f'Unet have {len(cross_kv_layer_names)} linear layer (related to text feature) need to optimize'
+    )
+
+    original_unet_state_dict = unet.state_dict()  # original state dict
+    concept_weights_dict = {}
+
+    # step 1: construct prompts for new concept -> extract input/target features
+    for concept, tuned_state_dict in zip(concept_list, unet_crosskv_list):
+
+        for layer_idx, layer_name in cross_kv_layer_names:
+
+            # merge params
+            original_params = original_unet_state_dict[layer_name]
+
+            # hard coded here: in unet, self/crosskv attention disable bias parameter
+            lora_down_name = layer_name.replace('to_k.weight', 'to_k.lora_down.weight').replace('to_v.weight', 'to_v.lora_down.weight')
+            lora_up_name = lora_down_name.replace('lora_down', 'lora_up')
+
+            alpha = concept['unet_alpha']
+
+            lora_down_params = tuned_state_dict[lora_down_name].to(device)
+            lora_up_params = tuned_state_dict[lora_up_name].to(device)
+
+            merge_params = original_params + alpha * lora_up_params @ lora_down_params
+
+            if layer_name not in concept_weights_dict:
+                concept_weights_dict[layer_name] = []
+
+            concept_weights_dict[layer_name].append(merge_params)
+
+
+    new_kv_weights = {}
+    # step 3: begin update model
+    for idx, (layer_idx, layer_name) in enumerate(cross_kv_layer_names):
+        Wnew = torch.stack(concept_weights_dict[layer_name])
+        Wnew = torch.mean(Wnew, dim = 0)
+        new_kv_weights[layer_name] = Wnew
+
+    return new_kv_weights
+
+
+def merge_text_encoder(concept_list, optimize_iters, new_concept_cfg,
+                       tokenizer, text_encoder, text_encoder_list, device):
+
+    LoRA_keys = []
+    for textenc_lora in text_encoder_list:
+        LoRA_keys += list(textenc_lora.keys())
+    LoRA_keys = set([
+        key.replace('.lora_down', '').replace('.lora_up', '')
+        for key in LoRA_keys
+    ])
+    text_encoder_layer_names = LoRA_keys
+
+    candidate_module_name = [
+        'q_proj', 'k_proj', 'v_proj', 'out_proj', 'fc1', 'fc2'
+    ]
+    candidate_module_name = [
+        name for name in candidate_module_name
+        if any([name in key for key in LoRA_keys])
+    ]
+
+    logging.info(f'text_encoder have {len(text_encoder_layer_names)} linear layer need to optimize')
+
+    global module_io_recoder, record_feature
+    hooker_handlers = []
+    for name, module in text_encoder.named_modules():
+        if any([item in name for item in candidate_module_name]):
+            hooker_handlers.append(module.register_forward_hook(hook=get_hooker(name)))
+
+    logging.info(f'add {len(hooker_handlers)} hooker to text_encoder')
+
+    original_state_dict = copy.deepcopy(text_encoder.state_dict())  # original state dict
+
+    new_concept_input_dict = {}
+    new_concept_output_dict = {}
+    concept_weights_dict = {}
+
+    for concept, lora_state_dict in zip(concept_list, text_encoder_list):
+        merged_state_dict = merge_lora_into_weight(
+            original_state_dict,
+            lora_state_dict,
+            text_encoder_layer_names,
+            model_type='text_encoder',
+            alpha=concept['text_encoder_alpha'],
+            device=device)
+        text_encoder.load_state_dict(merged_state_dict)  # load merged parameters
+        # we use different model to compute new concept feature
+        for layer_name in text_encoder_layer_names:
+            if layer_name not in concept_weights_dict:
+                concept_weights_dict[layer_name] = []
+            concept_weights_dict[layer_name].append(merged_state_dict[layer_name])
+
+    new_text_encoder_weights = {}
+    # step 3: begin update model
+    for idx, layer_name in enumerate(text_encoder_layer_names):
+        Wnew = torch.stack(concept_weights_dict[layer_name])
+        Wnew = torch.mean(Wnew, dim = 0)
+        new_text_encoder_weights[layer_name] = Wnew
+
+    logging.info(f'remove {len(hooker_handlers)} hooker from text_encoder')
+
+    # remove forward hooker
+    for hook_handle in hooker_handlers:
+        hook_handle.remove()
+
+    return new_text_encoder_weights
+
+
+@torch.no_grad()
+def decode_to_latents(concept_prompt, new_concept_cfg, tokenizer, text_encoder,
+                      unet, test_scheduler, num_inference_steps, device,
+                      record_nums, batch_size):
+
+    concept_prompt = bind_concept_prompt([concept_prompt], new_concept_cfg)
+    text_embeddings = get_text_feature(
+        concept_prompt,
+        tokenizer,
+        text_encoder,
+        device,
+        return_type='full_embedding').unsqueeze(0)
+
+    text_embeddings = text_embeddings.repeat((batch_size, 1, 1, 1))
+
+    # sd 1.x
+    height = 512
+    width = 512
+
+    latents = torch.randn((batch_size, unet.in_channels, height // 8, width // 8), )
+    latents = latents.to(device, dtype=text_embeddings.dtype)
+
+    test_scheduler.set_timesteps(num_inference_steps)
+    latents = latents * test_scheduler.init_noise_sigma
+
+    global record_feature
+    step = (test_scheduler.timesteps.size(0)) // record_nums
+    record_timestep = test_scheduler.timesteps[torch.arange(0, test_scheduler.timesteps.size(0), step=step)[:record_nums]]
+
+    for t in tqdm(test_scheduler.timesteps):
+
+        if t in record_timestep:
+            record_feature = True
+        else:
+            record_feature = False
+
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latent_model_input = latents
+        latent_model_input = test_scheduler.scale_model_input(latent_model_input, t)
+
+        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = test_scheduler.step(noise_pred, t, latents).prev_sample
+
+    return latents, text_embeddings
+
+
+def merge_spatial_attention(concept_list, optimize_iters, new_concept_cfg, tokenizer, text_encoder, unet, unet_spatial_attn_list, test_scheduler, device):
+    LoRA_keys = []
+    for unet_lora in unet_spatial_attn_list:
+        LoRA_keys += list(unet_lora.keys())
+    LoRA_keys = set([
+        key.replace('.lora_down', '').replace('.lora_up', '')
+        for key in LoRA_keys
+    ])
+    spatial_attention_layer_names = LoRA_keys
+
+    candidate_module_name = [
+        'attn2.to_q', 'attn2.to_out.0', 'attn1.to_q', 'attn1.to_k',
+        'attn1.to_v', 'attn1.to_out.0', 'ff.net.2', 'ff.net.0.proj',
+        'proj_out', 'proj_in'
+    ]
+    candidate_module_name = [
+        name for name in candidate_module_name
+        if any([name in key for key in LoRA_keys])
+    ]
+
+    logging.info(
+        f'unet have {len(spatial_attention_layer_names)} linear layer need to optimize'
+    )
+
+    global module_io_recoder
+    hooker_handlers = []
+    for name, module in unet.named_modules():
+        if any([x in name for x in candidate_module_name]):
+            hooker_handlers.append(
+                module.register_forward_hook(hook=get_hooker(name)))
+
+    logging.info(f'add {len(hooker_handlers)} hooker to unet')
+
+    original_state_dict = copy.deepcopy(unet.state_dict())  # original state dict
+    revise_edlora_unet_attention_forward(unet)
+
+    concept_weights_dict = {}
+
+    for concept, tuned_state_dict in zip(concept_list, unet_spatial_attn_list):
+        # set unet
+        module_io_recoder = {}  # reinit module io recorder
+
+        merged_state_dict = merge_lora_into_weight(
+            original_state_dict,
+            tuned_state_dict,
+            spatial_attention_layer_names,
+            model_type='unet',
+            alpha=concept['unet_alpha'],
+            device=device)
+        unet.load_state_dict(merged_state_dict)  # load merged parameters
+
+        concept_name = concept['concept_name']
+        concept_prompt = TEMPLATE_SIMPLE.format(concept_name)
+
+
+        for layer_name in spatial_attention_layer_names:
+            if layer_name not in concept_weights_dict:
+                concept_weights_dict[layer_name] = []
+
+            concept_weights_dict[layer_name].append(merged_state_dict[layer_name])
+
+    new_spatial_attention_weights = {}
+    # step 5: begin update model
+    for idx, layer_name in enumerate(spatial_attention_layer_names):
+        Wnew = torch.stack(concept_weights_dict[layer_name])
+        Wnew = torch.mean(Wnew, dim = 0)
+        new_spatial_attention_weights[layer_name] = Wnew
+
+    logging.info(f'remove {len(hooker_handlers)} hooker from unet')
+
+    for hook_handle in hooker_handlers:
+        hook_handle.remove()
+
+    return new_spatial_attention_weights
+
+
+def compose_concepts(concept_cfg, optimize_textenc_iters, optimize_unet_iters, pretrained_model_path, save_path, suffix, device):
+    logging.info('------Step 1: load stable diffusion checkpoint------')
+    pipe, train_scheduler, test_scheduler = init_stable_diffusion(pretrained_model_path, device)
+    tokenizer, text_encoder, unet, vae = pipe.tokenizer, pipe.text_encoder, pipe.unet, pipe.vae
+    for param in itertools.chain(text_encoder.parameters(), unet.parameters(), vae.parameters()):
+        param.requires_grad = False
+
+    logging.info('------Step 2: load new concepts checkpoints------')
+    embedding_list, text_encoder_list, unet_crosskv_list, unet_spatial_attn_list, concept_list = parse_new_concepts(concept_cfg)
+
+    # step 1: inplace add new concept to tokenizer and embedding layers of text encoder
+    if any([item is not None for item in embedding_list]):
+        logging.info('------Step 3: merge token embedding------')
+        _, new_concept_cfg = merge_new_concepts_(embedding_list, concept_list, tokenizer, text_encoder)
+    else:
+        _, new_concept_cfg = {}, {}
+        logging.info('------Step 3: no new embedding, skip merging token embedding------')
+
+    # step 2: construct reparameterized text_encoder
+    if any([item is not None for item in text_encoder_list]):
+        logging.info('------Step 4: merge text encoder------')
+        new_text_encoder_weights = merge_text_encoder(
+            concept_list, optimize_textenc_iters, new_concept_cfg, tokenizer,
+            text_encoder, text_encoder_list, device)
+        
+        # update the merged state_dict in text_encoder
+        text_encoder_state_dict = text_encoder.state_dict()
+        text_encoder_state_dict.update(new_text_encoder_weights)
+        text_encoder.load_state_dict(text_encoder_state_dict)
+    else:
+        new_text_encoder_weights = {}
+        logging.info('------Step 4: no new text encoder, skip merging text encoder------')
+
+    
+    # step 3: merge unet (k,v in crosskv-attention) params, since they only receive input from text-encoder
+
+    if any([item is not None for item in unet_crosskv_list]):
+        logging.info('------Step 5: merge kv of cross-attention in unet------')
+        new_kv_weights = merge_kv_in_cross_attention(
+            concept_list, optimize_textenc_iters, new_concept_cfg,
+            tokenizer, text_encoder, unet, unet_crosskv_list, device)
+        # update the merged state_dict in kv of crosskv-attention in Unet
+        unet_state_dict = unet.state_dict()
+        unet_state_dict.update(new_kv_weights)
+        unet.load_state_dict(unet_state_dict)
+    else:
+        new_kv_weights = {}
+        logging.info('------Step 5: no new kv of cross-attention in unet, skip merging kv------')
+
+    # step 4: merge unet (q,k,v in self-attention, q in crosskv-attention)
+    if any([item is not None for item in unet_spatial_attn_list]):
+        logging.info('------Step 6: merge spatial attention (q in cross-attention, qkv in self-attention) in unet------')
+        new_spatial_attention_weights = merge_spatial_attention(
+            concept_list, optimize_unet_iters, new_concept_cfg, tokenizer,
+            text_encoder, unet, unet_spatial_attn_list, test_scheduler, device)
+        unet_state_dict = unet.state_dict()
+        unet_state_dict.update(new_spatial_attention_weights)
+        unet.load_state_dict(unet_state_dict)
+    else:
+        new_spatial_attention_weights = {}
+        logging.info('------Step 6: no new spatial-attention in unet, skip merging spatial attention------')
+
+    checkpoint_save_path = f'{save_path}/combined_model_{suffix}'
+    pipe.save_pretrained(checkpoint_save_path)
+    with open(os.path.join(checkpoint_save_path, 'new_concept_cfg.json'), 'w') as json_file:
+        json.dump(new_concept_cfg, json_file)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser('', add_help=False)
+    parser.add_argument('--concept_cfg', help='json file for multi-concept', required=True, type=str)
+    parser.add_argument('--save_path', help='folder name to save optimized weights', required=True, type=str)
+    parser.add_argument('--suffix', help='suffix name', default='base', type=str)
+    parser.add_argument('--pretrained_models', required=True, type=str)
+    parser.add_argument('--optimize_unet_iters', default=50, type=int)
+    parser.add_argument('--optimize_textenc_iters', default=500, type=int)
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    # s1: set logger
+    exp_dir = f'{args.save_path}'
+    os.makedirs(exp_dir, exist_ok=True)
+    log_file = f'{exp_dir}/combined_model_{args.suffix}.log'
+    set_logger(log_file=log_file)
+    logging.info(args)
+
+    compose_concepts(args.concept_cfg,
+                     args.optimize_textenc_iters,
+                     args.optimize_unet_iters,
+                     args.pretrained_models,
+                     args.save_path,
+                     args.suffix,
+                     device='cuda')