hcpdiff 0.9.0__py3-none-any.whl → 2.1__py3-none-any.whl

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

hcpdiff/models/plugin.py

@@ -1,348 +0,0 @@
-"""
-plugin.py
-====================
-    :Name:        model plugin
-    :Author:      Dong Ziyi
-    :Affiliation: HCP Lab, SYSU
-    :Created:     10/03/2023
-    :Licence:     Apache-2.0
-"""
-
-import weakref
-import re
-from typing import Tuple, List, Dict, Any, Iterable
-
-import torch
-from torch import nn
-
-from hcpdiff.utils.net_utils import split_module_name
-
-class BasePluginBlock(nn.Module):
-    def __init__(self, name: str):
-        super().__init__()
-        self.name = name
-
-    def forward(self, host: nn.Module, fea_in: Tuple[torch.Tensor], fea_out: torch.Tensor):
-        return fea_out
-
-    def remove(self):
-        pass
-
-    def feed_input_data(self, data):
-        self.input_data = data
-
-    def register_input_feeder_to(self, host_model):
-        if not hasattr(host_model, 'input_feeder'):
-            host_model.input_feeder = []
-        host_model.input_feeder.append(self.feed_input_data)
-
-    def set_hyper_params(self, **kwargs):
-        for k, v in kwargs.items():
-            setattr(self, k, v)
-
-    @staticmethod
-    def extract_state_without_plugin(model: nn.Module, trainable=False):
-        trainable_keys = {k for k, v in model.named_parameters() if v.requires_grad}
-        plugin_names = {k for k, v in model.named_modules() if isinstance(v, BasePluginBlock)}
-        model_sd = {}
-        for k, v in model.state_dict().items():
-            if (not trainable) or k in trainable_keys:
-                for name in plugin_names:
-                    if k.startswith(name):
-                        break
-                else:
-                    model_sd[k] = v
-        return model_sd
-
-    def get_trainable_parameters(self) -> Iterable[nn.Parameter]:
-        return self.parameters()
-
-class WrapablePlugin:
-    wrapable_classes = ()
-
-    @classmethod
-    def wrap_layer(cls, name: str, layer: nn.Module, **kwargs):
-        plugin = cls(name, layer, **kwargs)
-        return plugin
-
-    @classmethod
-    def named_modules_with_exclude(cls, self, memo = None, prefix: str = '', remove_duplicate: bool = True,
-                                   exclude_key=None, exclude_classes=tuple()):
-
-        if memo is None:
-            memo = set()
-        if self not in memo:
-            if remove_duplicate:
-                memo.add(self)
-            if (exclude_key is None or not re.search(exclude_key, prefix)) and not isinstance(self, exclude_classes):
-                yield prefix, self
-                for name, module in self._modules.items():
-                    if module is None:
-                        continue
-                    submodule_prefix = prefix + ('.' if prefix else '') + name
-                    for m in cls.named_modules_with_exclude(module, memo, submodule_prefix, remove_duplicate, exclude_key, exclude_classes):
-                        yield m
-
-    @classmethod
-    def wrap_model(cls, name: str, host: nn.Module, exclude_key=None, exclude_classes=tuple(), **kwargs):  # -> Dict[str, SinglePluginBlock]:
-        '''
-        parent_block and other args required in __init__ will be put into kwargs, compatible with multiple models.
-        '''
-        plugin_block_dict = {}
-        if isinstance(host, cls.wrapable_classes):
-            plugin_block_dict[''] = cls.wrap_layer(name, host, **kwargs)
-        else:
-            named_modules = {layer_name:layer for layer_name, layer in cls.named_modules_with_exclude(
-                host, exclude_key=exclude_key, exclude_classes=exclude_classes)}
-            for layer_name, layer in named_modules.items():
-                if isinstance(layer, cls.wrapable_classes):
-                    # For plugins that need parent_block
-                    if 'parent_block' in kwargs:
-                        parent_name, host_name = split_module_name(layer_name)
-                        kwargs['parent_block'] = named_modules[parent_name]
-                        kwargs['host_name'] = host_name
-                    plugin_block_dict[layer_name] = cls.wrap_layer(name, layer, **kwargs)
-        return plugin_block_dict
-
-class SinglePluginBlock(BasePluginBlock, WrapablePlugin):
-
-    def __init__(self, name: str, host: nn.Module, hook_param=None, host_model=None):
-        super().__init__(name)
-        self.host = weakref.ref(host)
-        setattr(host, name, self)
-
-        if hook_param is None:
-            self.hook_handle = host.register_forward_hook(self.layer_hook)
-        else:  # hook for model parameters
-            self.backup = getattr(host, hook_param)
-            self.target = hook_param
-            self.handle_pre = host.register_forward_pre_hook(self.pre_hook)
-            self.handle_post = host.register_forward_hook(self.post_hook)
-
-    def layer_hook(self, host, fea_in: Tuple[torch.Tensor], fea_out: torch.Tensor):
-        return self(fea_in, fea_out)
-
-    def pre_hook(self, host, fea_in: torch.Tensor):
-        host.weight_restored = False
-        host_param = getattr(host, self.target)
-        delattr(host, self.target)
-        setattr(host, self.target, self(host_param))
-        return fea_in
-
-    def post_hook(self, host, fea_int, fea_out):
-        if not getattr(host, 'weight_restored', False):
-            setattr(host, self.target, self.backup)
-            host.weight_restored = True
-
-    def remove(self):
-        host = self.host()
-        delattr(host, self.name)
-        if hasattr(self, 'hook_handle'):
-            self.hook_handle.remove()
-        else:
-            self.handle_pre.remove()
-            self.handle_post.remove()
-
-class PluginBlock(BasePluginBlock):
-    def __init__(self, name, from_layer: Dict[str, Any], to_layer: Dict[str, Any], host_model=None):
-        super().__init__(name)
-        self.host_from = weakref.ref(from_layer['layer'])
-        self.host_to = weakref.ref(to_layer['layer'])
-        setattr(from_layer['layer'], name, self)
-
-        if from_layer['pre_hook']:
-            self.hook_handle_from = from_layer['layer'].register_forward_pre_hook(lambda host, fea_in:self.from_layer_hook(host, fea_in, None))
-        else:
-            self.hook_handle_from = from_layer['layer'].register_forward_hook(
-                lambda host, fea_in, fea_out:self.from_layer_hook(host, fea_in, fea_out))
-        if to_layer['pre_hook']:
-            self.hook_handle_to = to_layer['layer'].register_forward_pre_hook(lambda host, fea_in:self.to_layer_hook(host, fea_in, None))
-        else:
-            self.hook_handle_to = to_layer['layer'].register_forward_hook(lambda host, fea_in, fea_out:self.to_layer_hook(host, fea_in, fea_out))
-
-    def from_layer_hook(self, host, fea_in: Tuple[torch.Tensor], fea_out: torch.Tensor):
-        self.feat_from = fea_in
-
-    def to_layer_hook(self, host, fea_in: Tuple[torch.Tensor], fea_out: torch.Tensor):
-        return self(self.feat_from, fea_in, fea_out)
-
-    def remove(self):
-        host_from = self.host_from()
-        delattr(host_from, self.name)
-        self.hook_handle_from.remove()
-        self.hook_handle_to.remove()
-
-class MultiPluginBlock(BasePluginBlock):
-    def __init__(self, name: str, from_layers: List[Dict[str, Any]], to_layers: List[Dict[str, Any]], host_model=None):
-        super().__init__(name)
-        assert host_model is not None
-        self.host_from = [weakref.ref(x['layer']) for x in from_layers]
-        self.host_to = [weakref.ref(x['layer']) for x in to_layers]
-        self.host_model = weakref.ref(host_model)
-        setattr(host_model, name, self)
-
-        self.feat_from = [None for _ in range(len(from_layers))]
-
-        self.hook_handle_from = []
-        self.hook_handle_to = []
-
-        for idx, layer in enumerate(from_layers):
-            if layer['pre_hook']:
-                handle_from = layer['layer'].register_forward_pre_hook(lambda host, fea_in, idx=idx:self.from_layer_hook(host, fea_in, None, idx))
-            else:
-                handle_from = layer['layer'].register_forward_hook(
-                    lambda host, fea_in, fea_out, idx=idx:self.from_layer_hook(host, fea_in, fea_out, idx))
-            self.hook_handle_from.append(handle_from)
-        for idx, layer in enumerate(to_layers):
-            if layer['pre_hook']:
-                handle_to = layer['layer'].register_forward_pre_hook(lambda host, fea_in, idx=idx:self.to_layer_hook(host, fea_in, None, idx))
-            else:
-                handle_to = layer['layer'].register_forward_hook(lambda host, fea_in, fea_out, idx=idx:self.to_layer_hook(host, fea_in, fea_out, idx))
-            self.hook_handle_to.append(handle_to)
-
-        self.record_count = 0
-
-    def from_layer_hook(self, host, fea_in: Tuple[torch.Tensor], fea_out: Tuple[torch.Tensor], idx: int):
-        self.feat_from[idx] = fea_in
-        self.record_count += 1
-        if self.record_count == len(self.feat_from):  # call forward when all feat is record
-            self.record_count = 0
-            self.feat_to = self(self.feat_from)
-
-    def to_layer_hook(self, host, fea_in: Tuple[torch.Tensor], fea_out: Tuple[torch.Tensor], idx: int):
-        return self.feat_to[idx]+fea_out
-
-    def remove(self):
-        host_model = self.host_model()
-        delattr(host_model, self.name)
-        for handle_from in self.hook_handle_from:
-            handle_from.remove()
-        for handle_to in self.hook_handle_to:
-            handle_to.remove()
-
-class PatchPluginContainer(nn.Module):
-    def __init__(self, host_name, host, parent_block):
-        super().__init__()
-        self._host = host
-        self.host_name = host_name
-        self.parent_block = weakref.ref(parent_block)
-        self.plugin_names = []
-
-        delattr(parent_block, host_name)
-        setattr(parent_block, host_name, self)
-
-    def add_plugin(self, name: str, plugin: 'PatchPluginBlock'):
-        setattr(self, name, plugin)
-        self.plugin_names.append(name)
-
-    def remove_plugin(self, name: str):
-        delattr(self, name)
-        self.plugin_names.remove(name)
-        if len(self.plugin_names) == 0:
-            self.remove()
-
-    def forward(self, *args, **kwargs):
-        for name, plugin in self:
-            args, kwargs = plugin.pre_forward(*args, **kwargs)
-        output = self._host(*args, **kwargs)
-        for name, plugin in self:
-            output = plugin.post_forward(output, *args, **kwargs)
-        return output
-
-    def remove(self):
-        parent_block = self.parent_block()
-        delattr(parent_block, self.host_name)
-        setattr(parent_block, self.host_name, self._host)
-
-    def __iter__(self):
-        for name in self.plugin_names:
-            yield name, self[name]
-
-    def __getitem__(self, name):
-        return getattr(self, name)
-
-class PatchPluginBlock(BasePluginBlock, WrapablePlugin):
-    container_cls = PatchPluginContainer
-
-    def __init__(self, name: str, host: nn.Module, host_model=None, parent_block: nn.Module = None, host_name: str = None):
-        super().__init__(name)
-        if isinstance(host, self.container_cls):
-            self.host = weakref.ref(host._host)
-        else:
-            self.host = weakref.ref(host)
-        self.parent_block = weakref.ref(parent_block)
-        self.host_name = host_name
-
-        container = self.get_container(host, host_name, parent_block)
-        container.add_plugin(name, self)
-        self.container = weakref.ref(container)
-
-    def pre_forward(self, *args, **kwargs):
-        return args, kwargs
-
-    def post_forward(self, output, *args, **kwargs):
-        return output
-
-    def remove(self):
-        container = self.container()
-        container.remove_plugin(self.name)
-
-    def get_container(self, host, host_name, parent_block):
-        if isinstance(host, self.container_cls):
-            return host
-        else:
-            return self.container_cls(host_name, host, parent_block)
-
-    @classmethod
-    def wrap_model(cls, name: str, host: nn.Module, exclude_key=None, exclude_classes=tuple(), **kwargs):  # -> Dict[str, SinglePluginBlock]:
-        '''
-        parent_block and other args required in __init__ will be put into kwargs, compatible with multiple models.
-        '''
-        plugin_block_dict = {}
-        if isinstance(host, cls.wrapable_classes):
-            plugin_block_dict[''] = cls.wrap_layer(name, host, **kwargs)
-        else:
-            named_modules = {layer_name:layer for layer_name, layer in cls.named_modules_with_exclude(
-                host, exclude_key=exclude_key or '_host', exclude_classes=exclude_classes)}
-            for layer_name, layer in named_modules.items():
-                if isinstance(layer, cls.wrapable_classes) or isinstance(layer, cls.container_cls):
-                    # For plugins that need parent_block
-                    if 'parent_block' in kwargs:
-                        parent_name, host_name = split_module_name(layer_name)
-                        kwargs['parent_block'] = named_modules[parent_name]
-                        kwargs['host_name'] = host_name
-                    plugin_block_dict[layer_name] = cls.wrap_layer(name, layer, **kwargs)
-        return plugin_block_dict
-
-class PluginGroup:
-    def __init__(self, plugin_dict: Dict[str, BasePluginBlock]):
-        self.plugin_dict = plugin_dict  # {host_model_path: plugin_object}
-
-    def __setitem__(self, k, v):
-        self.plugin_dict[k] = v
-
-    def __getitem__(self, k):
-        return self.plugin_dict[k]
-
-    @property
-    def plugin_name(self):
-        if self.empty():
-            return None
-        return next(iter(self.plugin_dict.values())).name
-
-    def remove(self):
-        for plugin in self.plugin_dict.values():
-            plugin.remove()
-
-    def state_dict(self, model=None):
-        if model is None:
-            return {f'{k}.___.{ks}':vs for k, v in self.plugin_dict.items() for ks, vs in v.state_dict().items()}
-        else:
-            sd_model = model.state_dict()
-            return {f'{k}.___.{ks}':sd_model[f'{k}.{v.name}.{ks}'] for k, v in self.plugin_dict.items() for ks, vs in v.state_dict().items()}
-
-    def state_keys_raw(self):
-        return [f'{k}.{v.name}.{ks}' for k, v in self.plugin_dict.items() for ks, vs in v.state_dict().items()]
-
-    def empty(self):
-        return len(self.plugin_dict) == 0

hcpdiff/models/wrapper.py

@@ -1,75 +0,0 @@
-from torch import nn
-import itertools
-from transformers import CLIPTextModel
-from hcpdiff.utils import pad_attn_bias
-
-class TEUnetWrapper(nn.Module):
-    def __init__(self, unet, TE, train_TE=False):
-        super().__init__()
-        self.unet = unet
-        self.TE = TE
-
-        self.train_TE = train_TE
-
-    def forward(self, prompt_ids, noisy_latents, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
-        input_all = dict(prompt_ids=prompt_ids, noisy_latents=noisy_latents, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask, **plugin_input)
-
-        if hasattr(self.TE, 'input_feeder'):
-            for feeder in self.TE.input_feeder:
-                feeder(input_all)
-        encoder_hidden_states = self.TE(prompt_ids, position_ids=position_ids, attention_mask=attn_mask, output_hidden_states=True)[0]  # Get the text embedding for conditioning
-
-        if attn_mask is not None:
-            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
-
-        input_all['encoder_hidden_states'] = encoder_hidden_states
-        if hasattr(self.unet, 'input_feeder'):
-            for feeder in self.unet.input_feeder:
-                feeder(input_all)
-        model_pred = self.unet(noisy_latents, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask).sample  # Predict the noise residual
-        return model_pred
-
-    def prepare(self, accelerator):
-        if self.train_TE:
-            return accelerator.prepare(self)
-        else:
-            self.unet = accelerator.prepare(self.unet)
-            return self
-
-    def enable_gradient_checkpointing(self):
-        def grad_ckpt_enable(m):
-            if hasattr(m, 'gradient_checkpointing'):
-                m.training = True
-
-        self.unet.enable_gradient_checkpointing()
-        if self.train_TE:
-            self.TE.gradient_checkpointing_enable()
-            self.apply(grad_ckpt_enable)
-        else:
-            self.unet.apply(grad_ckpt_enable)
-
-    def trainable_parameters(self):
-        if self.train_TE:
-            return itertools.chain(self.unet.parameters(), self.TE.parameters())
-        else:
-            return self.unet.parameters()
-
-class SDXLTEUnetWrapper(TEUnetWrapper):
-    def forward(self, prompt_ids, noisy_latents, timesteps, attn_mask=None, position_ids=None, crop_info=None, plugin_input={}, **kwargs):
-        input_all = dict(prompt_ids=prompt_ids, noisy_latents=noisy_latents, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask, **plugin_input)
-
-        if hasattr(self.TE, 'input_feeder'):
-            for feeder in self.TE.input_feeder:
-                feeder(input_all)
-        encoder_hidden_states, pooled_output = self.TE(prompt_ids, position_ids=position_ids, attention_mask=attn_mask, output_hidden_states=True)  # Get the text embedding for conditioning
-
-        added_cond_kwargs = {"text_embeds":pooled_output[-1], "time_ids":crop_info}
-        if attn_mask is not None:
-            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
-
-        input_all['encoder_hidden_states'] = encoder_hidden_states
-        if hasattr(self.unet, 'input_feeder'):
-            for feeder in self.unet.input_feeder:
-                feeder(input_all)
-        model_pred = self.unet(noisy_latents, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask, added_cond_kwargs=added_cond_kwargs).sample  # Predict the noise residual
-        return model_pred

hcpdiff/noise/__init__.py

@@ -1,3 +0,0 @@
-from .noise_base import NoiseBase
-from .pyramid_noise import PyramidNoiseScheduler
-from .zero_terminal import ZeroTerminalScheduler

hcpdiff/noise/noise_base.py

@@ -1,16 +0,0 @@
-
-class NoiseBase:
-    def __init__(self, base_scheduler):
-        self.base_scheduler = base_scheduler
-
-    def __getattr__(self, item):
-        try:
-            return super(NoiseBase, self).__getattr__(item)
-        except:
-            return getattr(self.base_scheduler, item)
-
-    def __setattr__(self, key, value):
-        if  hasattr(super(), 'base_scheduler') and hasattr(self.base_scheduler, key):
-            setattr(self.base_scheduler, key, value)
-        else:
-            super(NoiseBase, self).__setattr__(key, value)

hcpdiff/noise/pyramid_noise.py

@@ -1,50 +0,0 @@
-import random
-
-import torch
-from torch.nn import functional as F
-from diffusers import SchedulerMixin
-
-from .noise_base import NoiseBase
-
-class PyramidNoiseScheduler(NoiseBase, SchedulerMixin):
-    def __init__(self, base_scheduler, level: int = 10, discount: float = 0.9, step_size: float = 2., resize_mode: str = 'bilinear'):
-        super().__init__(base_scheduler)
-        self.level = level
-        self.step_size = step_size
-        self.resize_mode = resize_mode
-        self.discount = discount
-
-    def add_noise(
-        self,
-        original_samples: torch.FloatTensor,
-        noise: torch.FloatTensor,
-        timesteps: torch.IntTensor,
-    ) -> torch.FloatTensor:
-        with torch.no_grad():
-            b, c, h, w = noise.shape
-            for i in range(1, self.level):
-                r = random.random()*2+self.step_size
-                wn, hn = max(1, int(w/(r**i))), max(1, int(h/(r**i)))
-                noise += F.interpolate(torch.randn(b, c, hn, wn).to(noise), (h, w), None, self.resize_mode)*(self.discount**i)
-                if wn == 1 or hn == 1:
-                    break
-            noise = noise/noise.std()
-        return self.base_scheduler.add_noise(original_samples, noise, timesteps)
-
-# if __name__ == '__main__':
-#     noise = torch.randn(1,3,512,512)
-#     level=10
-#     discount=0.6
-#     b, c, h, w = noise.shape
-#     for i in range(level):
-#         r = random.random() * 2 + 2
-#         wn, hn = max(1, int(w / (r ** i))), max(1, int(h / (r ** i)))
-#         noise += F.interpolate(torch.randn(b, c, wn, hn).to(noise), (w, h), None, 'bilinear') * discount ** i
-#         if wn == 1 or hn == 1:
-#             break
-#     noise = noise / noise.std()
-#
-#     from matplotlib import pyplot as plt
-#     plt.figure()
-#     plt.imshow(noise[0].permute(1,2,0))
-#     plt.show()