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

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)

hcpdiff/loggers/wandb_logger.py

@@ -1,31 +0,0 @@
-from typing import Dict, Any
-
-import os
-import wandb
-from PIL import Image
-
-from .base_logger import BaseLogger
-
-
-class WanDBLogger(BaseLogger):
-    def __init__(self, exp_dir, out_path=None, enable_log_image=False, project='hcp-diffusion', 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
-            wandb.init(project=project, name=os.path.basename(exp_dir))
-            wandb.save(os.path.join(exp_dir, 'cfg.yaml'), base_path=exp_dir)
-        else:
-            self.writer = None
-            self.disable()
-
-    def _info(self, info):
-        pass
-
-    def _log(self, datas: Dict[str, Any], step: int = 0):
-        log_dict = {'step': step}
-        for k, v in datas.items():
-            if len(v['data']) == 1:
-                log_dict[k] = v['data'][0]
-        wandb.log(log_dict)
-
-    def _log_image(self, imgs: Dict[str, Image.Image], step: int = 0):
-        wandb.log({next(iter(imgs.keys())): list(imgs.values())}, step=step)

hcpdiff/loggers/webui_logger.py

@@ -1,9 +0,0 @@
-from typing import Dict, Any
-
-from loguru import logger
-
-from .cli_logger import CLILogger
-
-class WebUILogger(CLILogger):
-    def _log(self, datas: Dict[str, Any], step: int = 0):
-        logger.info('this progress steps:'+', '.join([f"{k} = {v['format'].format(*v['data'])}" for k, v in datas.items()]))

hcpdiff/loss/min_snr_loss.py

@@ -1,52 +0,0 @@
-import torch
-from diffusers import SchedulerMixin
-from torch import nn
-
-class MinSNRLoss(nn.MSELoss):
-    need_timesteps = True
-
-    def __init__(self, size_average=None, reduce=None, reduction: str = 'none', gamma=1.,
-                 noise_scheduler: SchedulerMixin = None, device='cuda:0', **kwargs):
-        super().__init__(size_average, reduce, reduction)
-        self.gamma = gamma
-
-        # calculate SNR
-        alphas_cumprod = noise_scheduler.alphas_cumprod
-        sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
-        sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0-alphas_cumprod)
-        self.alpha = sqrt_alphas_cumprod.to(device)
-        self.sigma = sqrt_one_minus_alphas_cumprod.to(device)
-        self.all_snr = ((self.alpha/self.sigma)**2).to(device)
-
-    def forward(self, input: torch.Tensor, target: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
-        loss = super(MinSNRLoss, self).forward(input, target)
-        snr = self.all_snr[timesteps[:loss.shape[0], ...].squeeze()]
-        snr_weight = (self.gamma/snr).clip(max=1.).float()
-        return loss*snr_weight.view(-1, 1, 1, 1)
-
-
-class SoftMinSNRLoss(MinSNRLoss):
-    # gamma=2
-
-    def forward(self, input: torch.Tensor, target: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
-        loss = super(MinSNRLoss, self).forward(input, target)
-        snr = self.all_snr[timesteps[:loss.shape[0], ...].squeeze()]
-        snr_weight = (self.gamma**3/(snr**2 + self.gamma**3)).float()
-        return loss*snr_weight.view(-1, 1, 1, 1)
-
-class KDiffMinSNRLoss(MinSNRLoss):
-
-    def forward(self, input: torch.Tensor, target: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
-        loss = super(MinSNRLoss, self).forward(input, target)
-        snr = self.all_snr[timesteps[:loss.shape[0], ...].squeeze()]
-        snr_weight = 4*(((self.gamma*snr)**2/(snr**2 + self.gamma**2)**2)).float()
-        return loss*snr_weight.view(-1, 1, 1, 1)
-
-class EDMLoss(MinSNRLoss):
-
-    def forward(self, input: torch.Tensor, target: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
-        loss = super(MinSNRLoss, self).forward(input, target)
-        sigma = self.sigma[timesteps[:loss.shape[0], ...].squeeze()]
-        snr = self.all_snr[timesteps[:loss.shape[0], ...].squeeze()]
-        snr_weight = ((sigma**2+self.gamma**2)/(snr*(sigma*self.gamma)**2)).float()
-        return loss*snr_weight.view(-1, 1, 1, 1)

hcpdiff/models/layers.py

@@ -1,81 +0,0 @@
-"""
-layers.py
-====================
-    :Name:        GroupLinear and other layers
-    :Author:      Dong Ziyi
-    :Affiliation: HCP Lab, SYSU
-    :Created:     09/04/2023
-    :Licence:     Apache-2.0
-"""
-
-import torch
-from torch import nn
-import math
-from einops import rearrange
-
-class GroupLinear(nn.Module):
-    def __init__(self, in_features: int, out_features: int, groups: int, bias: bool = True,
-                 device=None, dtype=None):
-        super().__init__()
-        assert in_features%groups == 0
-        assert out_features%groups == 0
-
-        factory_kwargs = {'device': device, 'dtype': dtype}
-
-        self.groups = groups
-        self.in_features = in_features
-        self.out_features = out_features
-
-        self.weight = nn.Parameter(torch.empty((groups, in_features//groups, out_features//groups), **factory_kwargs))
-        if bias:
-            self.bias = nn.Parameter(torch.empty(groups, 1, out_features//groups, **factory_kwargs))
-        else:
-            self.register_parameter('bias', None)
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
-        # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
-        # https://github.com/pytorch/pytorch/issues/57109
-        self.kaiming_uniform_group(self.weight, a=math.sqrt(5))
-        if self.bias is not None:
-            fan_in, _ = self._calculate_fan_in_and_fan_out(self.weight)
-            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
-            nn.init.uniform_(self.bias, -bound, bound)
-
-    @staticmethod
-    def _calculate_fan_in_and_fan_out(tensor):
-        receptive_field_size = 1
-        num_input_fmaps = tensor.size(-2)
-        num_output_fmaps = tensor.size(-1)
-        fan_in = num_input_fmaps * receptive_field_size
-        fan_out = num_output_fmaps * receptive_field_size
-
-        return fan_in, fan_out
-
-    @staticmethod
-    def kaiming_uniform_group(tensor: torch.Tensor, a: float = 0, mode: str = 'fan_in', nonlinearity: str = 'leaky_relu') -> torch.Tensor:
-        def _calculate_correct_fan(tensor, mode):
-            mode = mode.lower()
-            valid_modes = ['fan_in', 'fan_out']
-            if mode not in valid_modes:
-                raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes))
-
-            fan_in, fan_out = GroupLinear._calculate_fan_in_and_fan_out(tensor)
-            return fan_in if mode == 'fan_in' else fan_out
-
-        fan = _calculate_correct_fan(tensor, mode)
-        gain = nn.init.calculate_gain(nonlinearity, a)
-        std = gain / math.sqrt(fan)
-        bound = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
-        with torch.no_grad():
-            return tensor.uniform_(-bound, bound)
-
-    def forward(self, x: torch.Tensor): # x: [G,B,L,C]
-        x = rearrange(x, '(g b) l c -> g (b l) c', g=self.num_groups)
-        if self.bias is not None:
-            out = torch.bmm(x, self.weight) + self.bias
-        else:
-            out = torch.bmm(x, self.weight)
-        out = rearrange(out, 'g (b l) c -> (g b) l c', b=B)
-        return out