utils_train.py

import shutil
import os
import os.path as osp
from pathlib import Path
from typing import Dict, List, Union

import argparse

import torch
from accelerate import Accelerator, DistributedType
from diffusers.models import AutoencoderKL


from library import train_util, chinese_sdxl_train_util

from removal.v1_2 import load_cfg, build_removal_model

def build_accelerator(args: argparse.Namespace, **kwargs) -> Accelerator:
    accelerator = train_util.prepare_accelerator(args, **kwargs)
    accelerator.print("prepare accelerator done")
    if accelerator.distributed_type == DistributedType.DEEPSPEED:
        # deepspeedの場合はtrain_micro_batch_size_per_gpuを設定しておく
        accelerator.state.deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu'] = 1
    return accelerator


def build_vae(model_cfg: Dict) -> AutoencoderKL:
    vae = AutoencoderKL.from_pretrained(model_cfg["vae"]['model_dir'])
    return vae


def build_models(args: argparse.Namespace, weight_dtype: str, accelerator: Accelerator) -> List[torch.nn.Module]:
    model_cfg = load_cfg(args.model_config_path)

    if args.pretrained_model_name_or_path: # interface for from_pretrained mathod
        model_cfg['vae']['model_dir'] = osp.join(args.pretrained_model_name_or_path, 'vae')
        model_cfg['model']['model_dir'] = osp.join(args.pretrained_model_name_or_path, 'unet')

    removal_model = build_removal_model(model_cfg, args.num_embeddings)
    vae = build_vae(model_cfg)

    accelerator.print(f"weight_dtype:{weight_dtype}")
    accelerator.print(f"vae:{vae.dtype}")
    if getattr(removal_model, 'unet', None):
        accelerator.print(f"unet:{removal_model.unet.dtype}")
    else:
        accelerator.print(f"diff_model:{removal_model.diff_model.dtype}")

    if args.reset_unet_parameters:
        accelerator.print("==> reset unet parameters")
        for name, layer in removal_model.unet.named_modules():
            if hasattr(layer, 'reset_parameters'):
                layer.reset_parameters()
                if accelerator.is_main_process:
                    accelerator.print(f'parameters reset: {name}')

    vae.requires_grad_(False).eval()
    vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
    vae.to(accelerator.device, dtype=vae_dtype)

    removal_model.to(accelerator.device, dtype=torch.float32)


    if accelerator.is_main_process:
        from pprint import pprint
        pprint("Model Config:")
        pprint(removal_model.diff_model.config)

    if args.gradient_checkpointing:
        removal_model.diff_model.enable_gradient_checkpointing()

    # set xformer/mem_eff_attn
    accelerator.print(f"Enable memory efficient attention, mem_eff_attn:{args.mem_eff_attn}, xformers:{args.xformers}")
    chinese_sdxl_train_util.set_diffusers_xformers_flag(removal_model.diff_model, True)
    chinese_sdxl_train_util.set_diffusers_xformers_flag(vae, True)

    return removal_model, vae




def build_dataloader(
    args: argparse.Namespace,
    # train_data_path: Union[str, Path], 
    dataset_class = None,
    accelerator: Accelerator = None
) -> List[Union[torch.utils.data.DataLoader, List[str]]]:
    def cycle(dl): yield from (data for _ in iter(int, 1) for data in dl)

    if args.data_config: # YAML dataset builder
        accelerator.print(f'[info]: build dataset from {args.data_config}')
        from omegaconf import OmegaConf, DictConfig
        d_cfg = OmegaConf.load(args.data_config).data
        train_data_path = d_cfg.path
        rand_mask_config = d_cfg.rand_mask_config
        use_rand_mask = d_cfg.use_rand_mask
        use_extra_fg_mask = d_cfg.use_extra_fg_mask
        ex_masks4pure_bg = d_cfg.extra_ann_files_4_PureBackTrain_2_RandMask
        train_jsons = train_data_path if not isinstance(train_data_path, str) else [train_data_path]
    else:
        accelerator.print(f'[info]: build dataset from proto args.')
        train_data_path = args.train_data_path
        rand_mask_config = args.rand_mask_config
        use_rand_mask = args.use_rand_mask
        use_extra_fg_mask = args.use_extra_fg_mask
        ex_masks4pure_bg = args.extra_ann_files_4_PureBackTrain_2_RandMask
        train_jsons = train_data_path if not isinstance(train_data_path, str) else [train_data_path.strip()]

    for i, json_file_path in enumerate(train_jsons):
            accelerator.print(f"[info]: ==> jsonl_idx:{i}, jsonl_path:{json_file_path}")

    train_dataset = dataset_class(
        ann_files=train_jsons,
        image_size=args.image_size,
        mask_config=rand_mask_config,
        extra_ann_files_4_PureBackTrain_2_RandMask=ex_masks4pure_bg,
        num_embeddings = args.num_embeddings,
        use_rand_mask=use_rand_mask,
        use_extra_fg_mask=use_extra_fg_mask,
        quiet=True # disable print on multi devices
    )

    accelerator.print(
        f'[info]: has {len(train_dataset.data_source_bg)} background task samples.')
    accelerator.print(
        f'[info]: has {len(train_dataset.data_source_fg)} foreground task samples.')
    accelerator.print(
        f'[info]: has {len(train_dataset.data_source)} total samples.')



    accelerator.print("[info]: copying train_jsons...")
    for train_json in train_jsons:
        dst_dir = osp.join(args.output_dir,"train_jsons")
        os.makedirs(dst_dir, exist_ok=True)
        dst_json = osp.join(dst_dir,osp.basename(train_json))
        if not os.path.exists(dst_json):
            shutil.copyfile(train_json, dst_json)

    batch_size, num_workers = args.batch_size, args.num_workers
    accelerator.print(f"[info]: batch_size is {batch_size}, num_workers is {num_workers}")
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        num_workers=num_workers, 
        shuffle=True, 
        drop_last=False,
        collate_fn=dataset_class.collate_fn
    )
    train_dataloader = accelerator.prepare(train_dataloader)
    return cycle(train_dataloader), train_jsons

def build_progress_bar(iterator, initial=0, 
    disable=False, desc='steps', mininterval=60, miniters=50, dynamic_ncols=False, 
    bar_format = "{l_bar}{bar:3}{r_bar}"):
    from tqdm import tqdm
    progress_bar = tqdm(
        iterator,  # range(len(train_dataloader)),
        initial=initial, 
        disable=disable, 
        mininterval=mininterval,
        miniters=miniters,
        bar_format=bar_format,
        dynamic_ncols=dynamic_ncols,
        desc=desc)
    return progress_bar



def save(model: torch.nn.Module, save_path: Union[str, Path], accelerator: Accelerator) -> None:
    accelerator.wait_for_everyone()
    removal_model_states = accelerator.unwrap_model(model).state_dict()
    if accelerator.is_main_process:
        os.makedirs(osp.dirname(save_path), exist_ok=True)
        torch.save(removal_model_states, save_path)
        accelerator.print(f"\n[info]: Model saved at: {save_path}\n")

    torch.cuda.empty_cache()


def common_arguments(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--no_half_vae", action="store_true", 
        help="do not use fp16/bf16 VAE in mixed precision (use float VAE)")
    parser.add_argument("--reset_unet_parameters", action="store_true", 
        help="reset unet parameters")
    parser.add_argument("--lognorm_t", action="store_true", help="whether lognorm timestep")

    parser.add_argument("--global_step", type=int, default=0, help="global_step")

    # data_cfg
    parser.add_argument("--train_data_path", type=str, nargs='+', default=None, 
        help="current train json data path, support multi paths split by space")
    parser.add_argument('--use_rand_mask', type=bool, default=True)
    parser.add_argument("--rand_mask_config", type=str, help="rand mask yaml",
        default="config/rand_mask_cfg/random_medium_512.yaml")
    parser.add_argument('--use_extra_fg_mask', type=bool, default=True)
    parser.add_argument("--extra_ann_files_4_PureBackTrain_2_RandMask", 
        type=str, default=None)
    parser.add_argument("--data_config", type=str, default=None)

    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_workers", type=int, default=8)

    parser.add_argument('--num_embeddings', type=int, default=20)
    parser.add_argument('--image_size', type=int, default=512)
    

    parser.add_argument('--model_config_path', type=str, default="")

    parser.add_argument('--cos_loss', action='store_true', default=False, help='whether use cosine similarity loss')

    parser.add_argument('--guidance_scale', type=float, default=1.0, help='class free guidance')

    parser.add_argument("--resume_from_ckpt", type=str, default="", help="resume from ckpt")