hcpdiff 0.9.1__py3-none-any.whl → 2.2__py3-none-any.whl

hcpdiff/data/pair_dataset.py

@@ -1,146 +0,0 @@
-"""
-pair_dataset.py
-====================
-    :Name:        text-image pair dataset
-    :Author:      Dong Ziyi
-    :Affiliation: HCP Lab, SYSU
-    :Created:     10/03/2023
-    :Licence:     Apache-2.0
-"""
-
-import os.path
-from argparse import Namespace
-
-import cv2
-import torch
-from PIL import Image
-from torch.utils.data import Dataset
-from tqdm.auto import tqdm
-from typing import Tuple
-
-from hcpdiff.utils.caption_tools import *
-from hcpdiff.utils.utils import get_file_name, get_file_ext
-from .bucket import BaseBucket
-from .source import DataSource, ComposeDataSource
-
-class TextImagePairDataset(Dataset):
-    """
-    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
-    It pre-processes the images and the tokenizes prompts.
-    """
-
-    def __init__(self, tokenizer, tokenizer_repeats: int = 1, att_mask_encode: bool = False,
-                 bucket: BaseBucket = None, source: Dict[str, DataSource] = None, return_path: bool = False,
-                 cache_path:str=None, encoder_attention_mask=False, **kwargs):
-        self.return_path = return_path
-
-        self.tokenizer = tokenizer
-        self.tokenizer_repeats = tokenizer_repeats
-        self.bucket: BaseBucket = bucket
-        self.att_mask_encode = att_mask_encode
-        self.source = ComposeDataSource(source)
-        self.latents = None  # Cache latents for faster training. Works only without image argumentations.
-        self.cache_path = cache_path
-        self.encoder_attention_mask = encoder_attention_mask
-
-    def load_data(self, path:str, data_source:DataSource, size:Tuple[int]):
-        image_dict = data_source.load_image(path)
-        image = image_dict['image']
-        att_mask = image_dict.get('att_mask', None)
-        if att_mask is None:
-            data, crop_coord = self.bucket.crop_resize({"img":image}, size)
-            image = data_source.procees_image(data['img'])  # resize to bucket size
-            att_mask = torch.ones((size[1]//8, size[0]//8))
-        else:
-            data, crop_coord = self.bucket.crop_resize({"img":image, "mask":att_mask}, size)
-            image = data_source.procees_image(data['img'])
-            att_mask = torch.tensor(cv2.resize(data['mask'], (size[0]//8, size[1]//8), interpolation=cv2.INTER_LINEAR))
-        return {'img':image, 'mask':att_mask}
-
-    @torch.no_grad()
-    def cache_latents(self, vae, weight_dtype, device, show_prog=True):
-        if self.cache_path and os.path.exists(self.cache_path):
-            self.latents = torch.load(self.cache_path)
-            return
-
-        self.latents = {}
-        self.bucket.rest(0)
-
-        for (path, data_source), size in tqdm(self.bucket, disable=not show_prog):
-            img_name = data_source.get_image_name(path)
-            if img_name not in self.latents:
-                data = self.load_data(path, data_source, size)
-                image = data['img'].unsqueeze(0).to(device, dtype=weight_dtype)
-                latents = vae.encode(image).latent_dist.sample().squeeze(0)
-                data['img'] = (latents*vae.config.scaling_factor).cpu()
-                self.latents[img_name] = data
-
-        if self.cache_path:
-            torch.save(self.latents, self.cache_path)
-
-    def __len__(self):
-        return len(self.bucket)
-
-    def __getitem__(self, index):
-        (path, data_source), size = self.bucket[index]
-        img_name = data_source.get_image_name(path)
-
-        if self.latents is None:
-            data = self.load_data(path, data_source, size)
-        else:
-            data = self.latents[img_name].copy()
-
-        prompt_ist = data_source.load_caption(img_name)
-
-        # tokenize Sp or (Sn, Sp)
-        tokens = self.tokenizer(prompt_ist, truncation=True, padding="max_length", return_tensors="pt",
-                                    max_length=self.tokenizer.model_max_length*self.tokenizer_repeats)
-        data['prompt'] = tokens.input_ids.squeeze()
-        if self.encoder_attention_mask and 'attention_mask' in tokens:
-            data['attn_mask'] = tokens.attention_mask.squeeze()
-        if 'position_ids' in tokens:
-            data['position_ids'] = tokens.position_ids.squeeze()
-
-        if self.return_path:
-            return data, path
-        else:
-            return data
-
-    @staticmethod
-    def collate_fn(batch):
-        '''
-        batch: [{img:tensor, prompt:str, ..., plugin_input:{...}},{}]
-        '''
-        has_plugin_input = 'plugin_input' in batch[0]
-        if has_plugin_input:
-            plugin_input = {k:[] for k in batch[0]['plugin_input'].keys()}
-
-        datas = {k:[] for k in batch[0].keys() if k != 'plugin_input' and k != 'prompt'}
-        sn_list, sp_list = [], []
-
-        for data in batch:
-            if has_plugin_input:
-                for k, v in data.pop('plugin_input').items():
-                    plugin_input[k].append(v)
-
-            prompt = data.pop('prompt')
-            if len(prompt.shape) == 2:
-                sn_list.append(prompt[0])
-                sp_list.append(prompt[1])
-            else:
-                sp_list.append(prompt)
-
-            for k, v in data.items():
-                datas[k].append(v)
-
-        for k, v in datas.items():
-            datas[k] = torch.stack(v)
-            if k == 'mask':
-                datas[k] = datas[k].unsqueeze(1)
-
-        sn_list += sp_list
-        datas['prompt'] = torch.stack(sn_list)
-        if has_plugin_input:
-            datas['plugin_input'] = {k:torch.stack(v) for k, v in plugin_input.items()}
-
-        return datas

hcpdiff/data/sampler.py

@@ -1,54 +0,0 @@
-import torch
-from torch.utils.data.distributed import DistributedSampler
-from typing import Iterator
-import platform
-import math
-
-class DistributedCycleSampler(DistributedSampler):
-    _cycle = 1000
-
-    def __iter__(self) -> Iterator:
-        def _iter():
-            while True:
-                if self.shuffle:
-                    # deterministically shuffle based on epoch and seed
-                    g = torch.Generator()
-                    g.manual_seed(self.seed + self.epoch)
-                    indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
-                else:
-                    indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
-
-                if not self.drop_last:
-                    # add extra samples to make it evenly divisible
-                    padding_size = self.total_size - len(indices)
-                    if padding_size <= len(indices):
-                        indices += indices[:padding_size]
-                    else:
-                        indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
-                else:
-                    # remove tail of data to make it evenly divisible.
-                    indices = indices[:self.total_size]
-                assert len(indices) == self.total_size
-
-                # subsample
-                indices = indices[self.rank:self.total_size:self.num_replicas]
-                assert len(indices) == self.num_samples
-
-                for idx in indices:
-                    yield idx
-                self.epoch+=1
-
-                if self.epoch>=self._cycle:
-                    break
-
-        return _iter()
-
-    def __len__(self) -> int:
-        return self.num_samples #*self._cycle
-
-def get_sampler():
-    # Fix DataLoader frequently reload bugs in windows
-    if platform.system().lower() == 'windows':
-        return DistributedCycleSampler
-    else:
-        return DistributedSampler

hcpdiff/data/source/base.py

@@ -1,30 +0,0 @@
-import os
-from typing import Dict, List, Tuple, Any
-
-class DataSource:
-    def __init__(self, img_root, repeat=1, **kwargs):
-        self.img_root = img_root
-        self.repeat = repeat
-
-    def get_image_list(self) -> List[Tuple[str, "DataSource"]]:
-        raise NotImplementedError()
-
-    def procees_image(self, image):
-        raise NotImplementedError()
-
-    def load_image(self, path) -> Dict[str, Any]:
-        raise NotImplementedError()
-
-    def get_image_name(self, path: str) -> str:
-        img_root, img_name = os.path.split(path)
-        return img_name.rsplit('.')[0]
-
-class ComposeDataSource(DataSource):
-    def __init__(self, source_dict: Dict[str, DataSource]):
-        self.source_dict = source_dict
-
-    def get_image_list(self) -> List[Tuple[str, DataSource]]:
-        img_list = []
-        for source in self.source_dict.values():
-            img_list.extend(source.get_image_list())
-        return img_list

hcpdiff/data/utils.py

@@ -1,80 +0,0 @@
-import cv2
-import numpy as np
-from PIL import Image
-from torchvision import transforms as T
-from torchvision.transforms import functional as F
-
-class DualRandomCrop:
-    def __init__(self, size):
-        self.size = size
-
-    def __call__(self, img):
-        crop_params = T.RandomCrop.get_params(img['img'], self.size)
-        img['img'] = F.crop(img['img'], *crop_params)
-        if "mask" in img:
-            img['mask'] = self.crop(img['mask'], *crop_params)
-        if "cond" in img:
-            img['cond'] = F.crop(img['cond'], *crop_params)
-        return img, crop_params[:2]
-
-    @staticmethod
-    def crop(img: np.ndarray, top: int, left: int, height: int, width: int) -> np.ndarray:
-        right = left+width
-        bottom = top+height
-        return img[top:bottom, left:right, ...]
-
-def resize_crop_fix(img, target_size, mask_interp=cv2.INTER_CUBIC):
-    w, h = img['img'].size
-    if w == target_size[0] and h == target_size[1]:
-        return img, [h,w,0,0,h,w]
-
-    ratio_img = w/h
-    if ratio_img>target_size[0]/target_size[1]:
-        new_size = (round(ratio_img*target_size[1]), target_size[1])
-        interp_type = Image.LANCZOS if h>target_size[1] else Image.BICUBIC
-    else:
-        new_size = (target_size[0], round(target_size[0]/ratio_img))
-        interp_type = Image.LANCZOS if w>target_size[0] else Image.BICUBIC
-    img['img'] = img['img'].resize(new_size, interp_type)
-    if "mask" in img:
-        img['mask'] = cv2.resize(img['mask'], new_size, interpolation=mask_interp)
-    if "cond" in img:
-        img['cond'] = img['cond'].resize(new_size, interp_type)
-
-    img, crop_coord = DualRandomCrop(target_size[::-1])(img)
-    return img, [*new_size, *crop_coord[::-1], *target_size]
-
-def pad_crop_fix(img, target_size):
-    w, h = img['img'].size
-    if w == target_size[0] and h == target_size[1]:
-        return img, (h,w,0,0,h,w)
-
-    pad_size = [0, 0, max(target_size[0]-w, 0), max(target_size[1]-h, 0)]
-    if pad_size[2]>0 or pad_size[3]>0:
-        img['img'] = F.pad(img['img'], pad_size)
-        if "mask" in img:
-            img['mask'] = np.pad(img['mask'], ((0, pad_size[3]), (0, pad_size[2])), 'constant', constant_values=(0, 0))
-        if "cond" in img:
-            img['cond'] = F.pad(img['cond'], pad_size)
-
-    if pad_size[2]>0 and pad_size[3]>0:
-        return img, (h,w,0,0,h,w)  # No need to crop
-    else:
-        img, crop_coord = DualRandomCrop(target_size[::-1])(img)
-        return img, crop_coord
-
-class CycleData():
-    def __init__(self, data_loader):
-        self.data_loader = data_loader
-
-    def __iter__(self):
-        self.epoch = 0
-
-        def cycle():
-            while True:
-                self.data_loader.dataset.bucket.rest(self.epoch)
-                for data in self.data_loader:
-                    yield data
-                self.epoch += 1
-
-        return cycle()

hcpdiff/deprecated/__init__.py

@@ -1 +0,0 @@
-from .lora_convert import convert_to_webui_maybe_old, convert_to_webui_xl_maybe_old

hcpdiff/deprecated/cfg_converter.py

@@ -1,81 +0,0 @@
-"""
-train_ac.py
-====================
-    :Name:        convert old cfg format to new format
-    :Author:      Dong Ziyi
-    :Affiliation: HCP Lab, SYSU
-    :Created:     10/03/2023
-    :Licence:     Apache-2.0
-"""
-
-from omegaconf import ListConfig, DictConfig, OmegaConf
-
-class DatasetCFGConverter:
-
-    def convert_source(self, cfg_source:DictConfig):
-        if '_target_' not in cfg_source:
-            cfg_source['_target_'] = 'hcpdiff.data.source.Text2ImageAttMapSource'
-
-        if 'tag_transforms' in cfg_source:
-            cfg_source['text_transforms'] = cfg_source.pop('tag_transforms')
-
-    def convert(self, cfg:DictConfig):
-        for dataset in cfg['data'].values():
-            for source in dataset['source'].values():
-                self.convert_source(source)
-        return cfg
-
-class TrainCFGConverter:
-    def __init__(self):
-        self.dataset_converter = DatasetCFGConverter()
-
-    def convert_model(self, cfg_model:DictConfig):
-        if 'ema_unet' in cfg_model and 'ema' not in cfg_model:
-            if cfg_model['ema_unet']==0: # no ema
-                cfg_model['ema'] = None
-            else:
-                cfg_model['ema'] = OmegaConf.create({
-                    '_target_': 'hcpdiff.utils.ema.ModelEMA',
-                    '_partial_': True,
-                    'decay_max': cfg_model['ema_unet'],
-                    'power': 0.85
-                })
-
-        if 'tokenizer' not in cfg_model:
-            cfg_model['tokenizer'] = None
-        if 'noise_scheduler' not in cfg_model:
-            cfg_model['noise_scheduler'] = None
-        if 'unet' not in cfg_model:
-            cfg_model['unet'] = None
-        if 'text_encoder' not in cfg_model:
-            cfg_model['text_encoder'] = None
-        if 'vae' not in cfg_model:
-            cfg_model['vae'] = None
-
-    def convert_loss(self, cfg_train:DictConfig):
-        if cfg_train['loss']['criterion']['_target_']=='hcpdiff.loss.MSELoss':
-            cfg_train['loss']['criterion']['_target_'] = 'torch.nn.MSELoss'
-
-    def convert(self, cfg:DictConfig):
-        self.convert_model(cfg['model'])
-        self.convert_loss(cfg['train'])
-
-        if 'previewer' not in cfg:
-            cfg['previewer'] = None
-
-        cfg = self.dataset_converter.convert(cfg)
-        return cfg
-
-class InferCFGConverter:
-
-    def convert(self, cfg:DictConfig):
-        if 'encoder_attention_mask' not in cfg:
-            cfg['encoder_attention_mask'] = False
-
-        if 'amp' not in cfg:
-            if cfg['dtype']=='amp':
-                cfg['dtype'] = 'fp32'
-                cfg['amp'] = True
-            else:
-                cfg['amp'] = False
-        return cfg

hcpdiff/deprecated/lora_convert.py

@@ -1,31 +0,0 @@
-
-def convert_to_webui_maybe_old(new_func):
-    def convert_to_webui_(self, state, prefix):
-        sd_covert = {}
-        for k, v in state.items():
-            # new lora format
-            if k.endswith('W_down'):
-                return new_func(self, state, prefix)
-
-            # old lora format
-            model_k, lora_k = k.split('.___.' if ('alpha' in k or 'scale' in k) else '.___.layer.', 1)
-            sd_covert[f"{prefix}{model_k.replace('.', '_')}.{lora_k}"] = v
-        return sd_covert
-    return convert_to_webui_
-
-def convert_to_webui_xl_maybe_old(new_func):
-    def convert_to_webui_xl_(self, state, prefix):
-        sd_convert = {}
-        for k, v in state.items():
-            # new lora format
-            if k.endswith('W_down'):
-                return new_func(self, state, prefix)
-
-            # old lora format
-            model_k, lora_k = k.split('.___.' if ('alpha' in k or 'scale' in k) else '.___.layer.', 1)
-            new_k = f"{prefix}{model_k.replace('.', '_')}.{lora_k}"
-            if 'clip' in new_k:
-                new_k = new_k.replace('_clip_B', '1') if 'clip_B' in new_k else new_k.replace('_clip_bigG', '2')
-            sd_convert[new_k] = v
-        return sd_convert
-    return convert_to_webui_xl_

hcpdiff/infer_workflow.py

@@ -1,57 +0,0 @@
-import argparse
-
-import torch
-import hydra
-from omegaconf import OmegaConf, DictConfig
-from easydict import EasyDict
-
-from hcpdiff.utils.utils import load_config_with_cli
-from .workflow import MemoryMixin
-from copy import deepcopy
-
-class WorkflowRunner:
-    def __init__(self, cfgs):
-        self.cfgs_raw = deepcopy(cfgs)
-        self.cfgs = OmegaConf.structured(cfgs, flags={"allow_objects": True})
-        OmegaConf.resolve(self.cfgs)
-        self.memory = EasyDict(**hydra.utils.instantiate(self.cfgs.memory))
-        self.attach_memory(self.cfgs)
-
-    def start(self):
-        prepare_actions = hydra.utils.instantiate(self.cfgs.prepare)
-        states = self.run(prepare_actions, {'cfgs': self.cfgs_raw})
-        actions = hydra.utils.instantiate(self.cfgs.actions)
-        states = self.run(actions, states)
-
-    def attach_memory(self, cfgs):
-        if OmegaConf.is_dict(cfgs):
-            if '_target_' in cfgs and cfgs['_target_'].endswith('.from_memory'):
-                cfgs._set_flag('allow_objects', True)
-                cfgs['memory'] = self.memory
-            else:
-                for v in cfgs.values():
-                    self.attach_memory(v)
-        elif OmegaConf.is_list(cfgs):
-            for v in cfgs:
-                self.attach_memory(v)
-
-    @torch.inference_mode()
-    def run(self, actions, states):
-        N_steps = len(actions)
-        for step, act in enumerate(actions):
-            print(f'[{step+1}/{N_steps}] action: {type(act).__name__}')
-            if isinstance(act, MemoryMixin):
-                states = act(memory=self.memory, **states)
-            else:
-                states = act(**states)
-            print(f'states: {", ".join(states.keys())}')
-        return states
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='HCP-Diffusion workflow')
-    parser.add_argument('--cfg', type=str, default='')
-    args, cfg_args = parser.parse_known_args()
-    cfgs = load_config_with_cli(args.cfg, args_list=cfg_args)  # skip --cfg
-
-    runner = WorkflowRunner(cfgs)
-    runner.start()

hcpdiff/loggers/__init__.py

@@ -1,13 +0,0 @@
-from .base_logger import BaseLogger, LoggerGroup
-from .cli_logger import CLILogger
-from .webui_logger import WebUILogger
-
-try:
-    from .tensorboard_logger import TBLogger
-except:
-    print('tensorboard is not available')
-
-try:
-    from .wandb_logger import WanDBLogger
-except:
-    print('wandb is not available')

hcpdiff/loggers/base_logger.py

@@ -1,76 +0,0 @@
-from typing import Dict, Any, List
-
-from PIL import Image
-
-from .preview import ImagePreviewer
-
-class BaseLogger:
-    def __init__(self, exp_dir, out_path, enable_log_image=False, log_step=10, image_log_step=200):
-        self.exp_dir = exp_dir
-        self.out_path = out_path
-        self.enable_log_image = enable_log_image
-        self.log_step = log_step
-        self.image_log_step = image_log_step
-        self.enable_log = True
-        self.previewer_list: List[ImagePreviewer] = []
-
-    def enable(self):
-        self.enable_log = True
-
-    def disable(self):
-        self.enable_log = False
-
-    def add_previewer(self, previewer: ImagePreviewer):
-        self.previewer_list.append(previewer)
-
-    def info(self, info):
-        if self.enable_log:
-            self._info(info)
-
-    def _info(self, info):
-        raise NotImplementedError()
-
-    def log(self, datas: Dict[str, Any], step: int = 0):
-        if self.enable_log and step%self.log_step == 0:
-            self._log(datas, step)
-
-    def _log(self, datas: Dict[str, Any], step: int = 0):
-        raise NotImplementedError()
-
-    def log_image(self, imgs: Dict[str, Image.Image], step: int = 0):
-        if self.enable_log and self.enable_log_image and step%self.image_log_step == 0:
-            self._log_image(imgs, step)
-
-    def _log_image(self, imgs: Dict[str, Image.Image], step: int = 0):
-        raise NotImplementedError()
-
-class LoggerGroup:
-    def __init__(self, logger_list: List[BaseLogger]):
-        self.logger_list = logger_list
-
-    def enable(self):
-        for logger in self.logger_list:
-            logger.enable()
-
-    def disable(self):
-        for logger in self.logger_list:
-            logger.disable()
-
-    def add_previewer(self, previewer):
-        for logger in self.logger_list:
-            logger.add_previewer(previewer)
-
-    def info(self, info):
-        for logger in self.logger_list:
-            logger.info(info)
-
-    def log(self, datas: Dict[str, Any], step: int = 0):
-        for logger in self.logger_list:
-            logger.log(datas, step)
-
-    def log_image(self, imgs: Dict[str, Image.Image], step: int = 0):
-        for logger in self.logger_list:
-            logger.log_image(imgs, step)
-
-    def __len__(self):
-        return len(self.logger_list)

hcpdiff/loggers/cli_logger.py

@@ -1,40 +0,0 @@
-import os
-from typing import Dict, Any
-
-from PIL import Image
-from loguru import logger
-
-from .base_logger import BaseLogger
-
-class CLILogger(BaseLogger):
-    def __init__(self, exp_dir, out_path, enable_log_image=False, log_step=10, image_log_step=200,
-                 img_log_dir='preview', img_ext='png', img_quality=95):
-        super().__init__(exp_dir, out_path, enable_log_image, log_step, image_log_step)
-        if exp_dir is not None:  # exp_dir is only available in local main process
-            logger.add(os.path.join(exp_dir, out_path))
-            if enable_log_image:
-                self.img_log_dir = os.path.join(exp_dir, img_log_dir)
-                os.makedirs(self.img_log_dir, exist_ok=True)
-            self.img_ext = img_ext
-            self.img_quality = img_quality
-        else:
-            self.disable()
-
-    def enable(self):
-        super(CLILogger, self).enable()
-        logger.enable("__main__")
-
-    def disable(self):
-        super(CLILogger, self).disable()
-        logger.disable("__main__")
-
-    def _info(self, info):
-        logger.info(info)
-
-    def _log(self, datas: Dict[str, Any], step: int = 0):
-        logger.info(', '.join([f"{k} = {v['format'].format(*v['data'])}" for k, v in datas.items()]))
-
-    def _log_image(self, imgs: Dict[str, Image.Image], step: int = 0):
-        logger.info(f'log {len(imgs)} images')
-        for name, img in imgs.items():
-            img.save(os.path.join(self.img_log_dir, f'{step}-{name}.{self.img_ext}'), quality=self.img_quality)

hcpdiff/loggers/preview/__init__.py

@@ -1 +0,0 @@
-from .image_previewer import ImagePreviewer