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

hcpdiff/ckpt_manager/ckpt_safetensor.py

@@ -1,64 +0,0 @@
-"""
-ckpt_safetensors.py
-====================
-    :Name:        save model with safetensors
-    :Author:      Dong Ziyi
-    :Affiliation: HCP Lab, SYSU
-    :Created:     8/04/2023
-    :Licence:     MIT
-"""
-
-import os
-import torch
-from safetensors import safe_open
-from safetensors.torch import save_file
-
-from .ckpt_pkl import CkptManagerPKL
-
-class CkptManagerSafe(CkptManagerPKL):
-
-    def _save_ckpt(self, sd_model, name=None, step=None, save_path=None):
-        if save_path is None:
-            save_path = os.path.join(self.save_dir, f"{name}-{step}.safetensors")
-        sd_unfold = self.unfold_dict(sd_model)
-        for k, v in sd_unfold.items():
-            if not v.is_contiguous():
-                sd_unfold[k] = v.contiguous()
-        save_file(sd_unfold, save_path)
-
-    def load_ckpt(self, ckpt_path, map_location='cpu'):
-        with safe_open(ckpt_path, framework="pt", device=map_location) as f:
-            sd_fold = self.fold_dict(f)
-        return sd_fold
-
-    @staticmethod
-    def unfold_dict(data, split_key=':'):
-        dict_unfold={}
-
-        def unfold(prefix, dict_fold):
-            for k,v in dict_fold.items():
-                k_new = k if prefix=='' else f'{prefix}{split_key}{k}'
-                if isinstance(v, dict):
-                    unfold(k_new, v)
-                elif isinstance(v, list) or isinstance(v, tuple):
-                    unfold(k_new, {i:d for i,d in enumerate(v)})
-                else:
-                    dict_unfold[k_new]=v
-
-        unfold('', data)
-        return dict_unfold
-
-    @staticmethod
-    def fold_dict(safe_f, split_key=':'):
-        dict_fold = {}
-
-        for k in safe_f.keys():
-            k_list = k.split(split_key)
-            dict_last = dict_fold
-            for item in k_list[:-1]:
-                if item not in dict_last:
-                    dict_last[item] = {}
-                dict_last = dict_last[item]
-            dict_last[k_list[-1]]=safe_f.get_tensor(k)
-
-        return dict_fold

hcpdiff/ckpt_manager/ckpt_webui.py

@@ -1,54 +0,0 @@
-from .base import CkptManagerBase
-import os
-from diffusers import StableDiffusionPipeline, UNet2DConditionModel
-from hcpdiff.models.plugin import BasePluginBlock
-from hcpdiff.tools.sd2diffusers import load_sd_ckpt, patch_method
-
-class CkptManagerWebui(CkptManagerBase):
-
-    def set_save_dir(self, save_dir, emb_dir=None):
-        os.makedirs(save_dir, exist_ok=True)
-        self.save_dir = save_dir
-        self.emb_dir = emb_dir
-
-    def save(self, step, unet, TE, lora_unet, lora_TE, all_plugin_unet, all_plugin_TE, embs, pipe: StableDiffusionPipeline, **kwargs):
-        def state_dict_unet(*args, model=unet, **kwargs):
-            plugin_names = {k for k, v in model.named_modules() if isinstance(v, BasePluginBlock)}
-            model_sd = {}
-            for k, v in model.state_dict_().items():
-                for name in plugin_names:
-                    if k.startswith(name):
-                        break
-                else:
-                    model_sd[k] = v
-            return model_sd
-        unet.state_dict_ = unet.state_dict
-        unet.state_dict = state_dict_unet
-
-        def state_dict_TE(*args, model=TE, **kwargs):
-            plugin_names = {k for k, v in model.named_modules() if isinstance(v, BasePluginBlock)}
-            model_sd = {}
-            for k, v in model.state_dict_().items():
-                for name in plugin_names:
-                    if k.startswith(name):
-                        break
-                else:
-                    model_sd[k] = v
-            return model_sd
-        TE.state_dict_ = TE.state_dict
-        TE.state_dict = state_dict_TE
-
-        pipe.save_pretrained(os.path.join(self.save_dir, f"model-{step}"), **kwargs)
-
-    @classmethod
-    def load(cls, pretrained_model, original_config_file, from_safetensors=False, device='cpu', ema=True, **kwargs) -> StableDiffusionPipeline:
-        patch_method()
-        pipe = load_sd_ckpt(
-            checkpoint_path=pretrained_model,
-            original_config_file=original_config_file,
-            extract_ema=ema,
-            scheduler_type='pndm',
-            from_safetensors=from_safetensors,
-            device=device,
-        )
-        return pipe

hcpdiff/data/bucket.py

@@ -1,358 +0,0 @@
-"""
-bucket.py
-====================
-    :Name:        aspect ratio bucket with k-means
-    :Author:      Dong Ziyi
-    :Affiliation: HCP Lab, SYSU
-    :Created:     10/03/2023
-    :Licence:     Apache-2.0
-"""
-
-import math
-import os.path
-import pickle
-from typing import List, Tuple, Union, Any
-
-import cv2
-import numpy as np
-from hcpdiff.utils.img_size_tool import types_support, get_image_size
-from hcpdiff.utils.utils import get_file_ext
-from .source import DataSource
-from loguru import logger
-from sklearn.cluster import KMeans
-from tqdm import tqdm
-from concurrent.futures import ThreadPoolExecutor
-
-from .utils import resize_crop_fix, pad_crop_fix
-
-class BaseBucket:
-    def __getitem__(self, idx):
-        '''
-        :return: (file name of image), (target image size)
-        '''
-        raise NotImplementedError()
-
-    def __len__(self):
-        raise NotImplementedError()
-
-    def build(self, bs: int, img_root_list: List[str]):
-        raise NotImplementedError()
-
-    def rest(self, epoch):
-        pass
-
-    def crop_resize(self, image, size, mask_interp=cv2.INTER_CUBIC) -> Tuple[Any, Tuple]:
-        return image, (*size, 0, 0, *size)
-
-class FixedBucket(BaseBucket):
-    def __init__(self, target_size: Union[Tuple[int, int], int] = 512, **kwargs):
-        self.target_size = (target_size, target_size) if isinstance(target_size, int) else target_size
-
-    def build(self, bs: int, file_names: List[Tuple[str, DataSource]]):
-        self.file_names = file_names
-
-    def crop_resize(self, image, size, mask_interp=cv2.INTER_CUBIC):
-        return resize_crop_fix(image, size, mask_interp=mask_interp)
-
-    def __getitem__(self, idx) -> Tuple[Tuple[str, DataSource], Tuple[int, int]]:
-        return self.file_names[idx], self.target_size
-
-    def __len__(self):
-        return len(self.file_names)
-
-class RatioBucket(BaseBucket):
-    def __init__(self, target_area: int = 640*640, step_size: int = 8, num_bucket: int = 10, pre_build_bucket: str = None):
-        self.target_area = target_area
-        self.step_size = step_size
-        self.num_bucket = num_bucket
-        self.pre_build_bucket = pre_build_bucket
-
-    def load_bucket(self, path):
-        with open(path, 'rb') as f:
-            data = pickle.load(f)
-        self.buckets = data['buckets']
-        self.size_buckets = data['size_buckets']
-        self.idx_bucket_map = data['idx_bucket_map']
-        self.data_len = data['data_len']
-
-    def save_bucket(self, path):
-        with open(path, 'wb') as f:
-            pickle.dump({
-                'buckets':self.buckets,
-                'size_buckets':self.size_buckets,
-                'idx_bucket_map':self.idx_bucket_map,
-                'data_len':self.data_len,
-            }, f)
-
-    def build_buckets_from_ratios(self):
-        logger.info('build buckets from ratios')
-        size_low = int(math.sqrt(self.target_area/self.ratio_max))
-        size_high = int(self.ratio_max*size_low)
-
-        # SD需要边长是8的倍数
-        size_low = (size_low//self.step_size)*self.step_size
-        size_high = (size_high//self.step_size)*self.step_size
-
-        data = []
-        for w in range(size_low, size_high+1, self.step_size):
-            for h in range(size_low, size_high+1, self.step_size):
-                data.append([w*h, np.log2(w/h), w, h])  # 对比例取对数，更符合人感知，宽高相反的可以对称分布。
-        data = np.array(data)
-
-        error_area = np.abs(data[:, 0]-self.target_area)
-        data_use = data[np.argsort(error_area)[:self.num_bucket*3], :]  # 取最小的num_bucket*3个
-
-        # 聚类，选出指定个数的bucket
-        kmeans = KMeans(n_clusters=self.num_bucket, random_state=3407).fit(data_use[:, 1].reshape(-1, 1))
-        labels = kmeans.labels_
-        self.buckets = []  # [bucket_id:[file_idx,...]]
-        ratios_log = []
-        self.size_buckets = []
-        for i in range(self.num_bucket):
-            map_idx = np.where(labels == i)[0]
-            m_idx = map_idx[np.argmin(np.abs(data_use[labels == i, 1]-np.median(data_use[labels == i, 1])))]
-            # self.buckets[wh_hash(*data_use[m_idx, 2:])]=[]
-            self.buckets.append([])
-            ratios_log.append(data_use[m_idx, 1])
-            self.size_buckets.append(data_use[m_idx, 2:].astype(int))
-        ratios_log = np.array(ratios_log)
-        self.size_buckets = np.array(self.size_buckets)
-
-        # fill buckets with images w,h
-        self.idx_bucket_map = np.empty(len(self.file_names), dtype=int)
-        for i, (file, source) in enumerate(self.file_names):
-            w, h = get_image_size(file)
-            bucket_id = np.abs(ratios_log-np.log2(w/h)).argmin()
-            self.buckets[bucket_id].append(i)
-            self.idx_bucket_map[i] = bucket_id
-        logger.info('buckets info: '+', '.join(f'size:{self.size_buckets[i]}, num:{len(b)}' for i, b in enumerate(self.buckets)))
-
-    def build_buckets_from_images(self):
-        logger.info('build buckets from images')
-
-        def get_ratio(data):
-            file, source = data
-            w, h = get_image_size(file)
-            ratio = np.log2(w/h)
-            return ratio
-
-        ratio_list = []
-        with ThreadPoolExecutor() as executor:
-            for ratio in tqdm(executor.map(get_ratio, self.file_names), desc='get image info', total=len(self.file_names)):
-                ratio_list.append(ratio)
-        ratio_list = np.array(ratio_list)
-
-        # 聚类，选出指定个数的bucket
-        kmeans = KMeans(n_clusters=self.num_bucket, random_state=3407, verbose=True, tol=1e-3).fit(ratio_list.reshape(-1, 1))
-        labels = kmeans.labels_
-        ratios = 2**kmeans.cluster_centers_.reshape(-1)
-
-        h_all = np.sqrt(self.target_area/ratios)
-        w_all = h_all*ratios
-
-        # SD需要边长是8的倍数
-        h_all = (np.round(h_all/self.step_size)*self.step_size).astype(int)
-        w_all = (np.round(w_all/self.step_size)*self.step_size).astype(int)
-        self.size_buckets = list(zip(w_all, h_all))
-        self.size_buckets = np.array(self.size_buckets)
-
-        self.buckets = []  # [bucket_id:[file_idx,...]]
-        self.idx_bucket_map = np.empty(len(self.file_names), dtype=int)
-        for bidx in range(self.num_bucket):
-            bnow = labels == bidx
-            self.buckets.append(np.where(bnow)[0].tolist())
-            self.idx_bucket_map[bnow] = bidx
-        logger.info('buckets info: '+', '.join(f'size:{self.size_buckets[i]}, num:{len(b)}' for i, b in enumerate(self.buckets)))
-
-    def build(self, bs: int, file_names: List[Tuple[str, DataSource]]):
-        '''
-        :param bs: batch_size * n_gpus * accumulation_step
-        :param img_root_list:
-        '''
-        self.file_names = file_names
-        self.bs = bs
-        if self.pre_build_bucket and os.path.exists(self.pre_build_bucket):
-            self.load_bucket(self.pre_build_bucket)
-            return
-
-        self._build()
-
-        rs = np.random.RandomState(42)
-        # make len(bucket)%bs==0
-        self.data_len = 0
-        for bidx, bucket in enumerate(self.buckets):
-            rest = len(bucket)%bs
-            if rest>0:
-                bucket.extend(rs.choice(bucket, bs-rest))
-            self.data_len += len(bucket)
-            self.buckets[bidx] = np.array(bucket)
-
-        if self.pre_build_bucket:
-            self.save_bucket(self.pre_build_bucket)
-
-    def rest(self, epoch):
-        rs = np.random.RandomState(42+epoch)
-        bucket_list = [x.copy() for x in self.buckets]
-        # shuffle inter bucket
-        for x in bucket_list:
-            rs.shuffle(x)
-
-        # shuffle of batches
-        bucket_list = np.hstack(bucket_list).reshape(-1, self.bs).astype(int)
-        rs.shuffle(bucket_list)
-
-        self.idx_bucket = bucket_list.reshape(-1)
-
-    def crop_resize(self, image, size, mask_interp=cv2.INTER_CUBIC):
-        return resize_crop_fix(image, size, mask_interp=mask_interp)
-
-    def __getitem__(self, idx):
-        file_idx = self.idx_bucket[idx]
-        bucket_idx = self.idx_bucket_map[file_idx]
-        return self.file_names[file_idx], self.size_buckets[bucket_idx]
-
-    def __len__(self):
-        return self.data_len
-
-    @classmethod
-    def from_ratios(cls, target_area: int = 640*640, step_size: int = 8, num_bucket: int = 10, ratio_max: float = 4,
-                    pre_build_bucket: str = None, **kwargs):
-        arb = cls(target_area, step_size, num_bucket, pre_build_bucket=pre_build_bucket)
-        arb.ratio_max = ratio_max
-        arb._build = arb.build_buckets_from_ratios
-        return arb
-
-    @classmethod
-    def from_files(cls, target_area: int = 640*640, step_size: int = 8, num_bucket: int = 10, pre_build_bucket: str = None, **kwargs):
-        arb = cls(target_area, step_size, num_bucket, pre_build_bucket=pre_build_bucket)
-        arb._build = arb.build_buckets_from_images
-        return arb
-
-class SizeBucket(RatioBucket):
-    def __init__(self, step_size: int = 8, num_bucket: int = 10, pre_build_bucket: str = None):
-        super().__init__(step_size=step_size, num_bucket=num_bucket, pre_build_bucket=pre_build_bucket)
-
-    def build_buckets_from_images(self):
-        '''
-        根据图像尺寸聚类，不会resize图像，只有剪裁和填充操作。
-        '''
-        logger.info('build buckets from images size')
-        size_list = []
-        for i, (file, source) in enumerate(self.file_names):
-            w, h = get_image_size(file)
-            size_list.append([w, h])
-        size_list = np.array(size_list)
-
-        # 聚类，选出指定个数的bucket
-        kmeans = KMeans(n_clusters=self.num_bucket, random_state=3407).fit(size_list)
-        labels = kmeans.labels_
-        size_buckets = kmeans.cluster_centers_
-
-        # SD需要边长是8的倍数
-        self.size_buckets = (np.round(size_buckets/self.step_size)*self.step_size).astype(int)
-
-        self.buckets = []  # [bucket_id:[file_idx,...]]
-        self.idx_bucket_map = np.empty(len(self.file_names), dtype=int)
-        for bidx in range(self.num_bucket):
-            bnow = labels == bidx
-            self.buckets.append(np.where(bnow)[0].tolist())
-            self.idx_bucket_map[bnow] = bidx
-        logger.info('buckets info: '+', '.join(f'size:{self.size_buckets[i]}, num:{len(b)}' for i, b in enumerate(self.buckets)))
-
-    def crop_resize(self, image, size):
-        return pad_crop_fix(image, size)
-
-    @classmethod
-    def from_files(cls, step_size: int = 8, num_bucket: int = 10, pre_build_bucket: str = None, **kwargs):
-        arb = cls(step_size, num_bucket, pre_build_bucket=pre_build_bucket)
-        arb._build = arb.build_buckets_from_images
-        return arb
-
-class RatioSizeBucket(RatioBucket):
-    def __init__(self, step_size: int = 8, num_bucket: int = 10, max_area:int=640*640, pre_build_bucket: str = None):
-        super().__init__(step_size=step_size, num_bucket=num_bucket, pre_build_bucket=pre_build_bucket)
-        self.max_area = max_area
-
-    def build_buckets_from_images(self):
-        '''
-        根据图像尺寸聚类，不会resize图像，只有剪裁和填充操作。
-        '''
-        logger.info('build buckets from images')
-        ratio_list = []
-        for i, (file, source) in enumerate(self.file_names):
-            w, h = get_image_size(file)
-            ratio = np.log2(w/h)
-            log_area = np.log2(min(w*h, self.max_area))
-            ratio_list.append([ratio, log_area])
-        ratio_list = np.array(ratio_list)
-
-        # 聚类，选出指定个数的bucket
-        kmeans = KMeans(n_clusters=self.num_bucket, random_state=3407).fit(ratio_list)
-        labels = kmeans.labels_
-        ratios = 2**kmeans.cluster_centers_[:, 0]
-        sizes = 2**kmeans.cluster_centers_[:, 1]
-
-        h_all = np.sqrt(sizes/ratios)
-        w_all = h_all*ratios
-
-        # SD需要边长是8的倍数
-        h_all = (np.round(h_all/self.step_size)*self.step_size).astype(int)
-        w_all = (np.round(w_all/self.step_size)*self.step_size).astype(int)
-        self.size_buckets = list(zip(w_all, h_all))
-        self.size_buckets = np.array(self.size_buckets)
-
-        self.buckets = []  # [bucket_id:[file_idx,...]]
-        self.idx_bucket_map = np.empty(len(self.file_names), dtype=int)
-        for bidx in range(self.num_bucket):
-            bnow = labels == bidx
-            self.buckets.append(np.where(bnow)[0].tolist())
-            self.idx_bucket_map[bnow] = bidx
-        logger.info('buckets info: '+', '.join(f'size:{self.size_buckets[i]}, num:{len(b)}' for i, b in enumerate(self.buckets)))
-
-    @classmethod
-    def from_files(cls, step_size: int = 8, num_bucket: int = 10, max_area:int=640*640, pre_build_bucket: str = None, **kwargs):
-        arb = cls(step_size, num_bucket, max_area=max_area, pre_build_bucket=pre_build_bucket)
-        arb._build = arb.build_buckets_from_images
-        return arb
-
-class LongEdgeBucket(RatioBucket):
-    def __init__(self, target_edge=640, step_size: int = 8, num_bucket: int = 10, pre_build_bucket: str = None):
-        super().__init__(step_size=step_size, num_bucket=num_bucket, pre_build_bucket=pre_build_bucket)
-        self.target_edge = target_edge
-
-    def build_buckets_from_images(self):
-        '''
-        根据图像尺寸聚类，不会resize图像，只有剪裁和填充操作。
-        '''
-        logger.info('build buckets from images size')
-        size_list = []
-        for i, (file, source) in enumerate(self.file_names):
-            w, h = get_image_size(file)
-            scale = self.target_edge/max(w, h)
-            size_list.append([round(w*scale), round(h*scale)])
-        size_list = np.array(size_list)
-
-        # 聚类，选出指定个数的bucket
-        kmeans = KMeans(n_clusters=self.num_bucket, random_state=3407, verbose=True).fit(size_list)
-        labels = kmeans.labels_
-        size_buckets = kmeans.cluster_centers_
-
-        # SD需要边长是8的倍数
-        self.size_buckets = (np.round(size_buckets/self.step_size)*self.step_size).astype(int)
-
-        self.buckets = []  # [bucket_id:[file_idx,...]]
-        self.idx_bucket_map = np.empty(len(self.file_names), dtype=int)
-        for bidx in range(self.num_bucket):
-            bnow = labels == bidx
-            self.buckets.append(np.where(bnow)[0].tolist())
-            self.idx_bucket_map[bnow] = bidx
-        logger.info('buckets info: '+', '.join(f'size:{self.size_buckets[i]}, num:{len(b)}' for i, b in enumerate(self.buckets)))
-
-    def crop_resize(self, image, size):
-        return resize_crop_fix(image, size)
-
-    @classmethod
-    def from_files(cls, target_edge, step_size: int = 8, num_bucket: int = 10, pre_build_bucket: str = None, **kwargs):
-        arb = cls(target_edge, step_size, num_bucket, pre_build_bucket=pre_build_bucket)
-        arb._build = arb.build_buckets_from_images
-        return arb

hcpdiff/data/caption_loader.py

@@ -1,80 +0,0 @@
-import json
-import os
-import glob
-import yaml
-from typing import Dict
-
-from loguru import logger
-from hcpdiff.utils.img_size_tool import types_support
-import os
-
-class BaseCaptionLoader:
-    def __init__(self, path):
-        self.path = path
-
-    def _load(self):
-        raise NotImplementedError
-        
-    def load(self):
-        retval = self._load()
-        logger.info(f'{len(retval)} record(s) loaded with {self.__class__.__name__}, from path {self.path!r}')
-        return retval
-
-    @staticmethod
-    def clean_ext(captions:Dict[str, str]):
-        def rm_ext(path):
-            name, ext = os.path.splitext(path)
-            if len(ext)>0 and ext[1:] in types_support:
-                return name
-            return path
-        return {rm_ext(k):v for k,v in captions.items()}
-
-class JsonCaptionLoader(BaseCaptionLoader):
-    def _load(self):
-        with open(self.path, 'r', encoding='utf-8') as f:
-            return self.clean_ext(json.loads(f.read()))
-
-class YamlCaptionLoader(BaseCaptionLoader):
-    def _load(self):
-        with open(self.path, 'r', encoding='utf-8') as f:
-            return self.clean_ext(yaml.load(f.read(), Loader=yaml.FullLoader))
-
-class TXTCaptionLoader(BaseCaptionLoader):
-    def _load(self):
-        txt_files = glob.glob(os.path.join(self.path, '*.txt'))
-        captions = {}
-        for file in txt_files:
-            with open(file, 'r', encoding='utf-8') as f:
-                captions[os.path.basename(file).split('.')[0]] = f.read().strip()
-        return captions
-
-def auto_caption_loader(path):
-    if os.path.isdir(path):
-        json_files = glob.glob(os.path.join(path, '*.json'))
-        if json_files:
-            return JsonCaptionLoader(json_files[0])
-
-        yaml_files = [
-            *glob.glob(os.path.join(path, '*.yaml')),
-            *glob.glob(os.path.join(path, '*.yml')),
-        ]
-        if yaml_files:
-            return YamlCaptionLoader(yaml_files[0])
-
-        txt_files = glob.glob(os.path.join(path, '*.txt'))
-        if txt_files:
-            return TXTCaptionLoader(path)
-
-        raise FileNotFoundError(f'Caption file not found in directory {path!r}.')
-
-    elif os.path.isfile(path):
-        _, ext = os.path.splitext(path)
-        if ext == '.json':
-            return JsonCaptionLoader(path)
-        elif ext in {'.yaml', '.yml'}:
-            return YamlCaptionLoader(path)
-        else:
-            raise FileNotFoundError(f'Unknown caption file {path!r}.')
-
-    else:
-        raise FileNotFoundError(f'Unknown caption file type {path!r}.')

hcpdiff/data/cond_dataset.py

@@ -1,40 +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 cv2
-import torch
-
-from .pair_dataset import TextImagePairDataset
-
-class TextImageCondPairDataset(TextImagePairDataset):
-    """
-    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 load_data(self, path, data_source, size):
-        image_dict = data_source.load_image(path)
-        image = image_dict['image']
-        att_mask = image_dict.get('att_mask', None)
-        img_cond = image_dict.get('cond', None)
-        if img_cond is None:
-            raise FileNotFoundError(f'{self.__class__} need the condition images!')
-
-        if att_mask is None:
-            data, crop_coord = self.bucket.crop_resize({"img":image, "cond":img_cond}, size)
-            image = data_source.procees_image(data['img'])  # resize to bucket size
-            img_cond = data_source.cond_transform(data['cond'])
-            att_mask = torch.ones((size[1]//8, size[0]//8))
-        else:
-            data, crop_coord = self.bucket.crop_resize({"img":image, "mask":att_mask, "cond":img_cond}, size)
-            image = data_source.procees_image(data['img'])
-            img_cond = data_source.cond_transform(data['cond'])
-            att_mask = torch.tensor(cv2.resize(data['mask'], (size[0]//8, size[1]//8), interpolation=cv2.INTER_LINEAR))
-        return {'img':image, 'mask':att_mask, 'plugin_input':{"cond":img_cond}}

hcpdiff/data/crop_info_dataset.py

@@ -1,40 +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
-"""
-
-from typing import Callable, Iterable, Dict
-from .bucket import BaseBucket
-import os.path
-
-import torch
-import cv2
-from .pair_dataset import TextImagePairDataset
-from hcpdiff.utils.utils import get_file_name
-from torchvision import transforms
-
-class CropInfoPairDataset(TextImagePairDataset):
-    """
-    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 load_data(self, path, data_source, size):
-        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))
-        crop_info = torch.tensor(crop_coord, dtype=torch.float)  # for sdxl
-        return {'img':image, 'mask':att_mask, 'crop_info':crop_info}

hcpdiff/data/data_processor.py

@@ -1,33 +0,0 @@
-import numpy as np
-import torch
-from PIL import Image
-from diffusers.utils import PIL_INTERPOLATION
-
-class ControlNetProcessor:
-    def __init__(self, image):
-        self.image_path = image
-
-    def prepare_cond_image(self, image, width, height, batch_size, device):
-        if not isinstance(image, torch.Tensor):
-            if isinstance(image, Image.Image):
-                image = [image]
-
-            if isinstance(image[0], Image.Image):
-                image = [
-                    np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image
-                ]
-                image = np.concatenate(image, axis=0)
-                image = np.array(image).astype(np.float32)/255.0
-                image = image.transpose(0, 3, 1, 2)
-                image = torch.from_numpy(image)
-            elif isinstance(image[0], torch.Tensor):
-                image = torch.cat(image, dim=0)
-
-        image = image.repeat_interleave(batch_size, dim=0)
-        image = image.to(device=device)
-
-        return image
-
-    def __call__(self, width, height, batch_size, device, dtype):
-        img = Image.open(self.image_path).convert('RGB')
-        return self.prepare_cond_image(img, width, height, batch_size, 'cuda').to(dtype=dtype)