hcpdiff 0.9.0__py3-none-any.whl → 2.1__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/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

hcpdiff/loggers/preview/image_previewer.py

@@ -1,149 +0,0 @@
-from contextlib import contextmanager
-from typing import List
-
-import hydra
-import torch
-from accelerate import infer_auto_device_map, dispatch_model
-from accelerate.hooks import remove_hook_from_module
-from diffusers import PNDMScheduler
-from torch.cuda.amp import autocast
-
-from hcpdiff.models import TokenizerHook
-from hcpdiff.utils.net_utils import to_cpu
-from hcpdiff.utils.utils import prepare_seed, load_config, size_to_int, int_to_size
-from hcpdiff.utils.utils import to_validate_file
-from hcpdiff.visualizer import Visualizer
-
-class ImagePreviewer(Visualizer):
-    def __init__(self, infer_cfg, exp_dir, te_hook,
-                 unet, TE, tokenizer, vae, save_cfg=False):
-        self.exp_dir = exp_dir
-        self.cfgs_raw = load_config(infer_cfg)
-        self.cfgs = hydra.utils.instantiate(self.cfgs_raw)
-        self.save_cfg = save_cfg
-        self.offload = 'offload' in self.cfgs and self.cfgs.offload is not None
-        self.dtype = self.dtype_dict[self.cfgs.dtype]
-
-        if getattr(self.cfgs.new_components, 'scheduler', None) is None:
-            scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule='scaled_linear')
-        else:
-            scheduler = self.cfgs.new_components.scheduler
-
-        pipe_cls = self.get_pipeline()
-        self.pipe = pipe_cls(vae=vae, text_encoder=TE, tokenizer=tokenizer, unet=unet, scheduler=scheduler, feature_extractor=None,
-                             safety_checker=None, requires_safety_checker=False)
-
-        self.token_ex = TokenizerHook(tokenizer)
-        self.te_hook = te_hook
-
-        if self.cfgs.seed is not None:
-            self.seeds = list(range(self.cfgs.seed, self.cfgs.seed+self.cfgs.num*self.cfgs.bs))
-        else:
-            self.seeds = [None]*(self.cfgs.num*self.cfgs.bs)
-
-    def build_vae_offload(self, offload_cfg):
-        vram = size_to_int(offload_cfg.max_VRAM)
-        if not offload_cfg.vae_cpu:
-            device_map = infer_auto_device_map(self.pipe.vae, max_memory={0:int_to_size(vram >> 5), "cpu":offload_cfg.max_RAM}, dtype=torch.float32)
-            self.pipe.vae = dispatch_model(self.pipe.vae, device_map)
-        else:
-            to_cpu(self.pipe.vae)
-            self.vae_decode_raw = self.pipe.vae.decode
-
-            def vae_decode_offload(latents, return_dict=True, decode_raw=self.pipe.vae.decode):
-                self.pipe.vae.to(dtype=torch.float32)
-                res = decode_raw(latents.cpu().to(dtype=torch.float32), return_dict=return_dict)
-                return res
-
-            self.pipe.vae.decode = vae_decode_offload
-
-            self.vae_encode_raw = self.pipe.vae.encode
-
-            def vae_encode_offload(x, return_dict=True, encode_raw=self.pipe.vae.encode):
-                self.pipe.vae.to(dtype=torch.float32)
-                res = encode_raw(x.cpu().to(dtype=torch.float32), return_dict=return_dict)
-                return res
-
-            self.pipe.vae.encode = vae_encode_offload
-
-    def remove_vae_offload(self, offload_cfg):
-        if not offload_cfg.vae_cpu:
-            remove_hook_from_module(self.pipe.vae, recurse=True)
-        else:
-            self.pipe.vae.encode = self.vae_encode_raw
-            self.pipe.vae.decode = self.vae_decode_raw
-
-    @contextmanager
-    def infer_optimize(self):
-        if getattr(self.cfgs, 'vae_optimize', None) is not None:
-            if self.cfgs.vae_optimize.tiling:
-                self.pipe.vae.enable_tiling()
-            if self.cfgs.vae_optimize.slicing:
-                self.pipe.vae.enable_slicing()
-        vae_device = self.pipe.vae.device
-        if self.offload:
-            self.build_vae_offload(self.cfgs.offload)
-        else:
-            self.pipe.vae.to(self.pipe.unet.device)
-
-        yield
-
-        if self.offload:
-            self.remove_vae_offload(self.cfgs.offload)
-        self.pipe.vae.to(vae_device)
-        self.pipe.vae.disable_tiling()
-        self.pipe.vae.disable_slicing()
-
-    def preview(self):
-        image_list, info_list = [], []
-        with self.infer_optimize():
-            for i in range(self.cfgs.num):
-                prompt = self.cfgs.prompt[i*self.cfgs.bs:(i+1)*self.cfgs.bs] if isinstance(self.cfgs.prompt, list) \
-                    else [self.cfgs.prompt]*self.cfgs.bs
-                negative_prompt = self.cfgs.neg_prompt[i*self.cfgs.bs:(i+1)*self.cfgs.bs] if isinstance(self.cfgs.neg_prompt, list) \
-                    else [self.cfgs.neg_prompt]*self.cfgs.bs
-                seeds = self.seeds[i*self.cfgs.bs:(i+1)*self.cfgs.bs]
-                images = self.vis_images(prompt=prompt, negative_prompt=negative_prompt, seeds=seeds,
-                                         **self.cfgs.infer_args)
-                for prompt_i, negative_prompt_i, seed in zip(prompt, negative_prompt, seeds):
-                    info_list.append({
-                        'prompt':prompt_i,
-                        'negative_prompt':negative_prompt_i,
-                        'seed':seed,
-                    })
-                image_list += images
-
-        return image_list, info_list
-
-    def preview_dict(self):
-        image_list, info_list = self.preview()
-        imgs = {f'{info["seed"]}-{to_validate_file(info["prompt"])}':img for img, info in zip(image_list, info_list)}
-        return imgs
-
-    @torch.no_grad()
-    def vis_images(self, prompt, negative_prompt='', seeds: List[int] = None, **kwargs):
-        G = prepare_seed(seeds or [None]*len(prompt))
-
-        ex_input_dict, pipe_input_dict = self.get_ex_input()
-        kwargs.update(pipe_input_dict)
-
-        mult_p, clean_text_p = self.token_ex.parse_attn_mult(prompt)
-        mult_n, clean_text_n = self.token_ex.parse_attn_mult(negative_prompt)
-        with autocast(enabled=self.cfgs.amp, dtype=self.dtype):
-            emb, pooled_output, attention_mask = self.te_hook.encode_prompt_to_emb(clean_text_n+clean_text_p)
-            if not self.cfgs.encoder_attention_mask:
-                attention_mask = None
-            emb_n, emb_p = emb.chunk(2)
-            emb_p = self.te_hook.mult_attn(emb_p, mult_p)
-            emb_n = self.te_hook.mult_attn(emb_n, mult_n)
-
-            if hasattr(self.pipe.unet, 'input_feeder'):
-                for feeder in self.pipe.unet.input_feeder:
-                    feeder(ex_input_dict)
-
-            if pooled_output is not None:
-                pooled_output = pooled_output[-1]
-
-            images = self.pipe(prompt_embeds=emb_p, negative_prompt_embeds=emb_n, callback=self.inter_callback, generator=G,
-                               pooled_output=pooled_output, encoder_attention_mask=attention_mask, **kwargs).images
-        return images

hcpdiff/loggers/tensorboard_logger.py

@@ -1,30 +0,0 @@
-import os
-from typing import Dict, Any
-
-import numpy as np
-from PIL import Image
-from torch.utils.tensorboard import SummaryWriter
-
-from .base_logger import BaseLogger
-
-
-class TBLogger(BaseLogger):
-    def __init__(self, exp_dir, out_path, enable_log_image=False, log_step=10, image_log_step=200):
-        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
-            self.writer = SummaryWriter(os.path.join(exp_dir, out_path))
-        else:
-            self.writer = None
-            self.disable()
-
-    def _info(self, info):
-        pass
-
-    def _log(self, datas: Dict[str, Any], step: int = 0):
-        for k, v in datas.items():
-            if len(v['data']) == 1:
-                self.writer.add_scalar(k, v['data'][0], global_step=step)
-
-    def _log_image(self, imgs: Dict[str, Image.Image], step: int = 0):
-        for name, img in imgs.items():
-            self.writer.add_image(f'img/{name}', np.array(img), dataformats='HWC', global_step=step)