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

hcpdiff/easy/cfg/t2i.py

@@ -0,0 +1,177 @@
+import torch
+from rainbowneko.infer.workflow import (Actions, PrepareAction, LoopAction, LoadModelAction)
+from rainbowneko.parser import neko_cfg, disable_neko_cfg
+from typing import Union, List
+
+from hcpdiff.ckpt_manager import HCPLoraLoader
+from hcpdiff.easy import Diffusers_SD, SD15_auto_loader, SDXL_auto_loader
+from hcpdiff.workflow import (BuildModelsAction, PrepareDiffusionAction, XformersEnableAction, VaeOptimizeAction, TextHookAction,
+                              AttnMultTextEncodeAction, SeedAction, MakeTimestepsAction, MakeLatentAction, DiffusionStepAction,
+                              time_iter, DecodeAction, SaveImageAction, LatentResizeAction)
+
+negative_prompt = 'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'
+
+## Easy config
+@neko_cfg
+def build_model(pretrained_model='ckpts/any5', noise_sampler=Diffusers_SD.dpmpp_2m_karras) -> Actions:
+    return Actions([
+        PrepareAction(device='cuda', dtype=torch.float16),
+        BuildModelsAction(
+            model_loader=SD15_auto_loader(
+                _partial_=True,
+                ckpt_path=pretrained_model,
+                noise_sampler=noise_sampler
+            )
+        ),
+    ])
+
+@neko_cfg
+def load_lora(info: List[List]) -> Actions:
+    lora_acts = []
+    for i, item in enumerate(info):
+        lora_unet = LoadModelAction(cfg={
+            f'lora_unet_{i}':HCPLoraLoader(
+                path=item[0],
+                state_prefix='denoiser.',
+                alpha=item[1],
+            )
+        }, key_map_in=('denoiser -> model', 'in_preview -> in_preview'))
+        lora_TE = LoadModelAction(cfg={
+            f'lora_unet_{i}':HCPLoraLoader(
+                path=item[0],
+                state_prefix='TE.',
+                alpha=item[1],
+            )
+        }, key_map_in=('TE -> model', 'in_preview -> in_preview'))
+
+        with disable_neko_cfg:
+            lora_acts.append(lora_unet)
+            lora_acts.append(lora_TE)
+
+    return Actions(lora_acts)
+
+@neko_cfg
+def optimize_model() -> Actions:
+    return Actions([
+        PrepareDiffusionAction(),
+        XformersEnableAction(),
+        VaeOptimizeAction(slicing=True),
+    ])
+
+@neko_cfg
+def text(prompt, negative_prompt=negative_prompt, bs=4) -> Actions:
+    return Actions([
+        TextHookAction(N_repeats=1, layer_skip=1),
+        AttnMultTextEncodeAction(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            bs=bs
+        ),
+    ])
+
+@neko_cfg
+def build_model_SDXL(pretrained_model='ckpts/any5', noise_sampler=Diffusers_SD.dpmpp_2m_karras) -> Actions:
+    return Actions([
+        PrepareAction(device='cuda', dtype=torch.float16),
+        ## Easy config
+        BuildModelsAction(
+            model_loader=SDXL_auto_loader(
+                _partial_=True,
+                ckpt_path=pretrained_model,
+                noise_sampler=noise_sampler
+            )
+        ),
+    ])
+
+@neko_cfg
+def text_SDXL(prompt, negative_prompt=negative_prompt, bs=4) -> Actions:
+    return Actions([
+        TextHookAction(N_repeats=1, layer_skip=1, TE_final_norm=False),
+        AttnMultTextEncodeAction(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            bs=bs
+        ),
+    ])
+
+@neko_cfg
+def config_diffusion(width=512, height=512, seed=None, N_steps=20, strength: float = None) -> Actions:
+    return Actions([
+        SeedAction(seed),
+        MakeTimestepsAction(N_steps=N_steps, strength=strength),
+        MakeLatentAction(width=width, height=height)
+    ])
+
+@neko_cfg
+def diffusion(guidance_scale=7.0) -> Actions:
+    return Actions([
+        LoopAction(
+            iterator=time_iter,
+            actions=[
+                DiffusionStepAction(guidance_scale=guidance_scale)
+            ]
+        )
+    ])
+
+@neko_cfg
+def decode(save_root='output_pipe/') -> Actions:
+    return Actions([
+        DecodeAction(),
+        SaveImageAction(save_root=save_root, image_type='png'),
+    ])
+
+@neko_cfg
+def resize(width=1024, height=1024):
+    return Actions([
+        LatentResizeAction(width=width, height=height)
+    ])
+
+@neko_cfg
+def SD15_t2i(pretrained_model, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4, width=512, height=512,
+             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+    return dict(workflow=Actions(actions=[
+        build_model(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        optimize_model(),
+        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
+        diffusion(guidance_scale=guidance_scale),
+        decode(save_root=save_root)
+    ]))
+
+@neko_cfg
+def SD15_t2i_lora(pretrained_model, lora_info, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4,
+                  width=512, height=512, seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+    return dict(workflow=Actions(actions=[
+        build_model(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        load_lora(info=lora_info),
+        optimize_model(),
+        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
+        diffusion(guidance_scale=guidance_scale),
+        decode(save_root=save_root)
+    ]))
+
+@neko_cfg
+def SDXL_t2i(pretrained_model, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4, width=1024, height=1024,
+             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+    return dict(workflow=Actions(actions=[
+        build_model_SDXL(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        optimize_model(),
+        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
+        diffusion(guidance_scale=guidance_scale),
+        decode(save_root=save_root)
+    ]))
+
+@neko_cfg
+def SDXL_t2i_lora(pretrained_model, lora_info, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4,
+                  width=1024, height=1024, seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+    return dict(workflow=Actions(actions=[
+        build_model_SDXL(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        load_lora(info=lora_info),
+        optimize_model(),
+        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
+        diffusion(guidance_scale=guidance_scale),
+        decode(save_root=save_root)
+    ]))

hcpdiff/easy/model/__init__.py

@@ -0,0 +1,2 @@
+from .loader import SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader
+from .cnet import ControlNet_SD15, make_controlnet_handler

hcpdiff/easy/model/cnet.py

@@ -0,0 +1,31 @@
+from hcpdiff.data.handler import ControlNetHandler, StableDiffusionHandler
+from hcpdiff.models import ControlNetPlugin
+from rainbowneko.data import SyncHandler
+from rainbowneko.parser import neko_cfg
+
+@neko_cfg
+def ControlNet_SD15(lr=1e-4):
+    return ControlNetPlugin(
+        _partial_=True,
+        lr=lr,
+        from_layers=[
+            'pre_hook:',
+            'pre_hook:conv_in',  # to make forward inside autocast
+        ],
+        to_layers=[
+            'down_blocks.0',
+            'down_blocks.1',
+            'down_blocks.2',
+            'down_blocks.3',
+            'mid_block',
+            'pre_hook:up_blocks.3.resnets.2',
+        ]
+    )
+
+@neko_cfg
+def make_controlnet_handler(bucket=None, encoder_attention_mask=False, erase=0.15, dropout=0.0, shuffle=0.0, word_names={}):
+    return SyncHandler(
+        diffusion=StableDiffusionHandler(bucket=bucket, encoder_attention_mask=encoder_attention_mask, erase=erase, dropout=dropout, shuffle=shuffle,
+                                         word_names=word_names),
+        cnet=ControlNetHandler(bucket=bucket)
+    )

hcpdiff/easy/model/loader.py

@@ -0,0 +1,79 @@
+import torch
+from hcpdiff.ckpt_manager import DiffusersSD15Format, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSD15Format, OfficialSDXLFormat
+from rainbowneko.ckpt_manager import NekoLoader, LocalCkptSource
+from hcpdiff.utils import auto_tokenizer_cls, auto_text_encoder_cls, get_pipe_name
+from hcpdiff.models.wrapper import SDXLWrapper, SD15Wrapper, PixArtWrapper
+from hcpdiff.models.compose import SDXLTextEncoder
+from diffusers import  StableDiffusionPipeline, StableDiffusionXLPipeline
+
+def SD15_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    try:
+        try_diffusers = StableDiffusionPipeline.load_config(ckpt_path)
+        loader = NekoLoader(
+            format=DiffusersSD15Format(),
+            source=LocalCkptSource(),
+        )
+    except EnvironmentError:
+        loader = NekoLoader(
+            format=OfficialSD15Format(),
+            source=LocalCkptSource(),
+        )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype, **kwargs)
+    return models
+
+def SDXL_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    try:
+        try_diffusers = StableDiffusionXLPipeline.load_config(ckpt_path)
+        loader = NekoLoader(
+            format=DiffusersSDXLFormat(),
+            source=LocalCkptSource(),
+        )
+    except EnvironmentError:
+        loader = NekoLoader(
+            format=OfficialSDXLFormat(),
+            source=LocalCkptSource(),
+        )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype, **kwargs)
+    return models
+
+def PixArt_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    loader = NekoLoader(
+        format=DiffusersPixArtFormat(),
+        source=LocalCkptSource(),
+    )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype, **kwargs)
+    return models
+
+def auto_load_wrapper(pretrained_model, denoiser=None, TE=None, vae=None, noise_sampler=None, tokenizer=None, revision=None,
+                      dtype=torch.float32, **kwargs):
+    if TE is not None:
+        text_encoder_cls = type(TE)
+    else:
+        text_encoder_cls = auto_text_encoder_cls(pretrained_model, revision)
+
+    pipe_name = get_pipe_name(pretrained_model)
+
+    if text_encoder_cls == SDXLTextEncoder:
+        wrapper_cls = SDXLWrapper
+        format = DiffusersSDXLFormat()
+    elif 'PixArt' in pipe_name:
+        wrapper_cls = PixArtWrapper
+        format = DiffusersPixArtFormat()
+    else:
+        wrapper_cls = SD15Wrapper
+        format = DiffusersSD15Format()
+
+    loader = NekoLoader(
+        format=format,
+        source=LocalCkptSource(),
+    )
+    models = loader.load(pretrained_model, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype)
+
+    return wrapper_cls.build_from_pretrained(models, **kwargs)

hcpdiff/easy/sampler.py

@@ -0,0 +1,46 @@
+from hcpdiff.diffusion.sampler import DiffusersSampler
+from diffusers import DPMSolverMultistepScheduler, DDIMScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler
+
+class Diffusers_SD:
+    dpmpp_2m = DiffusersSampler(
+        DPMSolverMultistepScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+            algorithm_type='dpmsolver++',
+        )
+    )
+
+    dpmpp_2m_karras = DiffusersSampler(
+        DPMSolverMultistepScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+            algorithm_type='dpmsolver++',
+            use_karras_sigmas=True,
+        )
+    )
+
+    ddim = DiffusersSampler(
+        DDIMScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+        )
+    )
+
+    euler = DiffusersSampler(
+        EulerDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+        )
+    )
+
+    euler_a = DiffusersSampler(
+        EulerAncestralDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+        )
+    )

hcpdiff/evaluate/__init__.py

@@ -0,0 +1 @@
+from .previewer import HCPPreviewer

hcpdiff/evaluate/previewer.py

@@ -0,0 +1,60 @@
+from pathlib import Path
+
+import torch
+from rainbowneko.evaluate.preview import WorkflowPreviewer
+from rainbowneko.utils import to_cuda
+
+from hcpdiff.models.wrapper import SD15Wrapper
+from accelerate.hooks import remove_hook_from_module
+
+class HCPPreviewer(WorkflowPreviewer):
+
+    @torch.no_grad()
+    def evaluate(self, step: int, model: SD15Wrapper, prefix='eval/'):
+        if step%self.interval != 0 or not self.trainer.is_local_main_process:
+            return
+
+        # record training layers
+        training_layers = [layer for layer in model.modules() if layer.training]
+
+        model.eval()
+        self.trainer.loggers.info(f'Preview')
+
+        N_repeats = model.text_enc_hook.N_repeats
+        clip_skip = model.text_enc_hook.clip_skip
+        clip_final_norm = model.text_enc_hook.clip_final_norm
+        use_attention_mask = model.text_enc_hook.use_attention_mask
+
+        preview_root = Path(self.trainer.exp_dir)/'imgs'
+        preview_root.mkdir(parents=True, exist_ok=True)
+
+        states = self.workflow_runner.run(model=model, in_preview=True, te_hook=model.text_enc_hook,
+                                          device=self.device, dtype=self.dtype, preview_root=preview_root, preview_step=step,
+                                          world_size=self.trainer.world_size, local_rank=self.trainer.local_rank,
+                                          emb_hook=self.trainer.cfgs.emb_pt.embedding_hook if self.trainer.pt_trainable else None)
+
+        # restore model states
+        if model.vae is not None:
+            model.vae.disable_tiling()
+            model.vae.disable_slicing()
+            remove_hook_from_module(model.vae, recurse=True)
+            if 'vae_encode_raw' in states:
+                model.vae.encode = states['vae_encode_raw']
+                model.vae.decode = states['vae_decode_raw']
+
+        if 'emb_hook' in states and not self.trainer.pt_trainable:
+            states['emb_hook'].remove()
+
+        if self.trainer.pt_trainable:
+            self.trainer.cfgs.emb_pt.embedding_hook.N_repeats = N_repeats
+
+        model.tokenizer.N_repeats = N_repeats
+        model.text_enc_hook.N_repeats = N_repeats
+        model.text_enc_hook.clip_skip = clip_skip
+        model.text_enc_hook.clip_final_norm = clip_final_norm
+        model.text_enc_hook.use_attention_mask = use_attention_mask
+        
+        to_cuda(model)
+
+        for layer in training_layers:
+            layer.train()

hcpdiff/loss/__init__.py

@@ -1 +1,4 @@
-from .min_snr_loss import MinSNRLoss, SoftMinSNRLoss, KDiffMinSNRLoss, EDMLoss
+from .weighting import MinSNRWeight, SNRWeight, EDMWeight, LossWeight
+from .ssim import SSIMLoss, MS_SSIMLoss
+from .gw import GWLoss
+from .base import DiffusionLossContainer

hcpdiff/loss/base.py

@@ -0,0 +1,41 @@
+from rainbowneko.train.loss import LossContainer
+from typing import Dict, Any
+from torch import Tensor
+
+class DiffusionLossContainer(LossContainer):
+    def __init__(self, loss, weight=1.0, key_map=None):
+        key_map = key_map or getattr(loss, '_key_map', None) or ('pred.model_pred -> 0', 'pred.target -> 1')
+        super().__init__(loss, weight, key_map)
+        self.target_type = getattr(loss, 'target_type', 'eps')
+
+    def get_target(self, pred_type, model_pred, x_0, noise, x_t, sigma, noise_sampler, **kwargs):
+        # Get target
+        if self.target_type == "eps":
+            target = noise
+        elif self.target_type == "x0":
+            target = x_0
+        elif self.target_type == "velocity":
+            target = noise_sampler.eps_to_velocity(noise, x_t, sigma)
+        else:
+            raise ValueError(f"Unsupport target_type {self.target_type}")
+
+        # TODO: put in wrapper
+        # # remove pred vars
+        # if model_pred.shape[1] == target.shape[1]*2:
+        #     model_pred, _ = model_pred.chunk(2, dim=1)
+
+        # Convert pred_type to target_type
+        if pred_type != self.target_type:
+            cvt_func = getattr(noise_sampler, f'{pred_type}_to_{self.target_type}', None)
+            if cvt_func is None:
+                raise ValueError(f"Unsupport pred_type {pred_type} with target_type {self.target_type}")
+            else:
+                model_pred = cvt_func(model_pred, x_t, sigma)
+        return model_pred, target
+    
+    def forward(self, pred:Dict[str,Any], inputs:Dict[str,Any]) -> Tensor:
+        model_pred, target = self.get_target(**pred)
+        pred['model_pred'] = model_pred
+        pred['target'] = target
+        loss = super().forward(pred, inputs) * self.weight # [B,*,*,*]
+        return loss.mean()

hcpdiff/loss/gw.py

@@ -0,0 +1,35 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class GWLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        sobel_x = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
+        sobel_y = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]
+        self.sobel_x = torch.FloatTensor(sobel_x)
+        self.sobel_y = torch.FloatTensor(sobel_y)
+        self.register_buffer('sobel_x', self.sobel_x)
+        self.register_buffer('sobel_y', self.sobel_y)
+
+    def forward(self, pred, target):
+        '''
+
+        :param pred: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,C,H,W]
+        '''
+        b, c, w, h = pred.shape
+
+        sobel_x = self.sobel_x.expand(c, 1, 3, 3).to(pred.device)
+        sobel_y = self.sobel_y.expand(c, 1, 3, 3).to(pred.device)
+        Ix1 = F.conv2d(pred, sobel_x, stride=1, padding=1, groups=c)
+        Ix2 = F.conv2d(target, sobel_x, stride=1, padding=1, groups=c)
+        Iy1 = F.conv2d(pred, sobel_y, stride=1, padding=1, groups=c)
+        Iy2 = F.conv2d(target, sobel_y, stride=1, padding=1, groups=c)
+
+        dx = torch.abs(Ix1 - Ix2)
+        dy = torch.abs(Iy1 - Iy2)
+        loss = (1 + 4 * dx) * (1 + 4 * dy) * torch.abs(pred - target)
+        return loss

hcpdiff/loss/ssim.py

@@ -0,0 +1,37 @@
+from pytorch_msssim import SSIM, MS_SSIM
+from torch.nn.modules.loss import _Loss
+import torch
+
+class SSIMLoss(_Loss):
+    target_type = 'x0'
+
+    def __init__(self, size_average=None, reduce=None, reduction: str = 'mean'):
+        super().__init__(size_average=size_average, reduce=reduce, reduction=reduction)
+        self.ssim = SSIM(data_range=1., size_average=False, channel=4)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        '''
+
+        :param input: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,1,1,1]
+        '''
+        input = (input+1)/2
+        target = (target+1)/2
+        return 1-self.ssim(input, target).view(-1,1,1,1)
+
+class MS_SSIMLoss(_Loss):
+    def __init__(self, size_average=None, reduce=None, reduction: str = 'mean'):
+        super().__init__(size_average=size_average, reduce=reduce, reduction=reduction)
+        self.ssim = MS_SSIM(data_range=1., size_average=False, channel=4)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        '''
+
+        :param input: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,1,1,1]
+        '''
+        input = (input+1)/2
+        target = (target+1)/2
+        return 1-self.ssim(input, target).view(-1,1,1,1)

hcpdiff/loss/vlb.py

@@ -0,0 +1,79 @@
+import torch
+from torch import nn
+import numpy as np
+
+class VLBLoss(nn.Module):
+    need_sigma = True
+    need_timesteps = True
+    need_sampler = True
+    var_pred = True
+
+    def __init__(self, loss, weight: float = 1.):
+        super().__init__()
+        self.loss = loss
+        self.weight = weight
+
+    def normal_kl(self, mean1, logvar1, mean2, logvar2):
+        """
+        Compute the KL divergence between two gaussians.
+        """
+
+        return 0.5*(-1.0+logvar2-logvar1+(logvar1-logvar2).exp()+((mean1-mean2)**2)*(-logvar2).exp())
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor, sigma, timesteps: torch.Tensor, x_t: torch.Tensor, sampler):
+        eps_pred, var_pred = input.chunk(2, dim=1)
+        x0_pred = sampler.eps_to_x0(eps_pred, x_t, sigma)
+
+        true_mean = sampler.sigma_scheduler.get_post_mean(timesteps, target, x_t)
+        true_logvar = sampler.sigma_scheduler.get_post_log_var(timesteps)
+
+        pred_mean = sampler.sigma_scheduler.get_post_mean(timesteps, x0_pred, x_t)
+        pred_logvar = sampler.sigma_scheduler.get_post_log_var(timesteps, x_t_var=var_pred)
+
+        kl = self.normal_kl(true_mean, true_logvar, pred_mean, pred_logvar)
+        kl = kl.mean(dim=(1,2,3))/np.log(2.0)
+
+        decoder_nll = -self.discretized_gaussian_log_likelihood(target, means=pred_mean, log_scales=0.5*pred_logvar)
+        assert decoder_nll.shape == target.shape
+        decoder_nll = decoder_nll.mean(dim=(1,2,3))/np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = torch.where((timesteps == 0), decoder_nll, kl)
+
+        return self.weight*output
+
+    def approx_standard_normal_cdf(self, x):
+        """
+        A fast approximation of the cumulative distribution function of the
+        standard normal.
+        """
+        return 0.5*(1.0+torch.tanh(np.sqrt(2.0/np.pi)*(x+0.044715*torch.pow(x, 3))))
+
+    def discretized_gaussian_log_likelihood(self, x, *, means, log_scales):
+        """
+        Compute the log-likelihood of a Gaussian distribution discretizing to a
+        given image.
+        :param x: the target images. It is assumed that this was uint8 values,
+                  rescaled to the range [-1, 1].
+        :param means: the Gaussian mean Tensor.
+        :param log_scales: the Gaussian log stddev Tensor.
+        :return: a tensor like x of log probabilities (in nats).
+        """
+        assert x.shape == means.shape == log_scales.shape
+        centered_x = x-means
+        inv_stdv = torch.exp(-log_scales)
+        plus_in = inv_stdv*(centered_x+1.0/255.0)
+        cdf_plus = self.approx_standard_normal_cdf(plus_in)
+        min_in = inv_stdv*(centered_x-1.0/255.0)
+        cdf_min = self.approx_standard_normal_cdf(min_in)
+        log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
+        log_one_minus_cdf_min = torch.log((1.0-cdf_min).clamp(min=1e-12))
+        cdf_delta = cdf_plus-cdf_min
+        log_probs = torch.where(
+            x<-0.999,
+            log_cdf_plus,
+            torch.where(x>0.999, log_one_minus_cdf_min, torch.log(cdf_delta.clamp(min=1e-12))),
+        )
+        assert log_probs.shape == x.shape
+        return log_probs

hcpdiff/loss/weighting.py

@@ -0,0 +1,66 @@
+from torch import nn
+
+from .base import DiffusionLossContainer
+
+class LossWeight(nn.Module):
+    def __init__(self, loss: DiffusionLossContainer):
+        super().__init__()
+        self.loss = loss
+
+    def get_weight(self, pred, inputs):
+        '''
+
+        :param input: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,1,1,1] or [B,C,H,W]
+        '''
+        raise NotImplementedError
+
+    def forward(self, pred, inputs):
+        '''
+        weight: [B,1,1,1] or [B,C,H,W]
+        loss: [B,*,*,*]
+        '''
+        return self.get_weight(pred, inputs)*self.loss(pred, inputs)
+
+class SNRWeight(LossWeight):
+    def get_weight(self, pred, inputs):
+        if self.loss.target_type == 'eps':
+            return 1
+        elif self.loss.target_type == "x0":
+            sigma = pred['sigma']
+            return (1./sigma**2).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+
+class MinSNRWeight(LossWeight):
+    def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
+        super().__init__(loss)
+        self.gamma = gamma
+
+    def get_weight(self, pred, inputs):
+        sigma = pred['sigma']
+        if self.loss.target_type == 'eps':
+            w_snr = (self.gamma*sigma**2).clip(max=1).float()
+        elif self.loss.target_type == "x0":
+            w_snr = (1/(sigma**2)).clip(max=self.gamma).float()
+        else:
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+
+        return w_snr.view(-1, 1, 1, 1)
+
+class EDMWeight(LossWeight):
+    def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
+        super().__init__(loss)
+        self.gamma = gamma
+
+    def get_weight(self, pred, inputs):
+        sigma = pred['sigma']
+        if self.loss.target_type == 'eps':
+            w_snr = ((sigma**2+self.gamma**2)/(self.gamma**2)).float()
+        elif self.loss.target_type == "x0":
+            w_snr = ((sigma**2+self.gamma**2)/((sigma*self.gamma)**2)).float()
+        else:
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+
+        return w_snr.view(-1, 1, 1, 1)

hcpdiff/models/__init__.py

@@ -1,4 +1,3 @@
-from .plugin import PluginBlock, PluginGroup, SinglePluginBlock, MultiPluginBlock, PatchPluginBlock
 # from .lora_base import LoraBlock, LoraGroup
 # from .lora_layers import lora_layer_map
 from .lora_base_patch import LoraBlock, LoraGroup
@@ -7,4 +6,5 @@ from .text_emb_ex import EmbeddingPTHook
 from .textencoder_ex import TEEXHook
 from .tokenizer_ex import TokenizerHook
 from .cfg_context import CFGContext, DreamArtistPTContext
-from .wrapper import TEUnetWrapper, SDXLTEUnetWrapper
+from .wrapper import SD15Wrapper, SDXLWrapper, PixArtWrapper, TEHookCFG
+from .controlnet import ControlNetPlugin