hcpdiff 2.2.1__py3-none-any.whl → 2.3__py3-none-any.whl

hcpdiff/easy/cfg/sdxl_train.py

@@ -1,15 +1,14 @@
 import torch
-from rainbowneko.ckpt_manager import ckpt_saver, NekoPluginSaver, LAYERS_TRAINABLE, SafeTensorFormat
-from rainbowneko.parser import CfgWDPluginParser, neko_cfg, CfgWDModelParser, disable_neko_cfg
-from rainbowneko.utils import ConstantLR
-
+from hcpdiff.ckpt_manager import LoraWebuiFormat
 from hcpdiff.easy import SDXL_auto_loader
 from hcpdiff.models import SDXLWrapper
 from hcpdiff.models.lora_layers_patch import LoraLayer
-from hcpdiff.ckpt_manager import LoraWebuiFormat
+from rainbowneko.ckpt_manager import ckpt_saver, NekoPluginSaver, LAYERS_TRAINABLE, SafeTensorFormat, NekoOptimizerSaver
+from rainbowneko.parser import CfgWDPluginParser, neko_cfg, CfgWDModelParser, disable_neko_cfg
+from rainbowneko.utils import ConstantLR
 
 @neko_cfg
-def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int = 500, lr: float = 1e-5,
+def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int = 500, save_optimizer=False, lr: float = 1e-5,
                     dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0, name: str = 'SDXL'):
     if low_vram:
         from bitsandbytes.optim import AdamW8bit
@@ -17,6 +16,17 @@ def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
     else:
         optimizer = torch.optim.AdamW(_partial_=True)
 
+    ckpt_saver_dict = dict(
+        SDXL=ckpt_saver(
+            ckpt_type='safetensors',
+            target_module='denoiser',
+            layers=LAYERS_TRAINABLE,
+        )
+    )
+
+    if save_optimizer:
+        ckpt_saver_dict['optimizer'] = NekoOptimizerSaver()
+
     from cfgs.train.py import train_base, tuning_base
 
     return dict(
@@ -30,13 +40,7 @@ def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
             )
         ], weight_decay=1e-2),
 
-        ckpt_saver=dict(
-            SDXL=ckpt_saver(
-                ckpt_type='safetensors',
-                target_module='denoiser',
-                layers=LAYERS_TRAINABLE,
-            )
-        ),
+        ckpt_saver=ckpt_saver_dict,
 
         train=dict(
             train_steps=train_steps,
@@ -64,9 +68,9 @@ def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
     )
 
 @neko_cfg
-def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int = 200, lr: float = 1e-4, rank: int = 4, alpha: float = None,
-                    with_conv: bool = False, dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0, name: str = 'SDXL',
-                    save_webui_format=False):
+def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int = 200, save_optimizer=False, lr: float = 1e-4, rank: int = 4,
+                    alpha: float = None, with_conv: bool = False, dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0,
+                    name: str = 'SDXL', save_webui_format=False):
     with disable_neko_cfg:
         if alpha is None:
             alpha = rank
@@ -97,6 +101,17 @@ def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
     else:
         lora_format = SafeTensorFormat()
 
+    ckpt_saver_dict = dict(
+        _replace_=True,
+        lora_unet=NekoPluginSaver(
+            format=lora_format,
+            target_plugin='lora1',
+        )
+    )
+
+    if save_optimizer:
+        ckpt_saver_dict['optimizer'] = NekoOptimizerSaver()
+
     from cfgs.train.py.examples import SD_FT
 
     return dict(
@@ -114,13 +129,7 @@ def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
             )
         ), weight_decay=0.1),
 
-        ckpt_saver=dict(
-            _replace_ = True,
-            lora_unet=NekoPluginSaver(
-                format=lora_format,
-                target_plugin='lora1',
-            )
-        ),
+        ckpt_saver=ckpt_saver_dict,
 
         train=dict(
             train_steps=train_steps,

hcpdiff/evaluate/__init__.py

@@ -1 +1,3 @@
-from .previewer import HCPPreviewer
+from .previewer import HCPPreviewer
+from .evaluator import HCPEvaluator
+from .metrics import CLIPScoreMetric

hcpdiff/evaluate/evaluator.py

@@ -0,0 +1,76 @@
+from pathlib import Path
+
+import torch
+from accelerate.hooks import remove_hook_from_module
+from rainbowneko.evaluate import WorkflowEvaluator, MetricGroup
+from rainbowneko.utils import to_cuda
+
+from hcpdiff.models.wrapper import SD15Wrapper
+
+class HCPEvaluator(WorkflowEvaluator):
+
+    @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)
+
+        # get metrics
+        metric = states['_metric']
+
+        v_metric = metric.finish(self.trainer.accelerator.gather, self.trainer.is_local_main_process)
+        if not isinstance(v_metric, dict):
+            v_metric = {'metric':v_metric}
+
+        log_data = {
+            "eval/Step":{
+                "format":"{}",
+                "data":[step],
+            }
+        }
+        log_data.update(MetricGroup.format(v_metric, prefix=prefix))
+        self.trainer.loggers.log(log_data, step, force=True)
+
+        # 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/evaluate/metrics/__init__.py

@@ -0,0 +1 @@
+from .clip_score import CLIPScoreMetric

hcpdiff/evaluate/metrics/clip_score.py

@@ -0,0 +1,23 @@
+from torchmetrics.multimodal.clip_score import CLIPScore, _clip_score_update
+from torch import Tensor
+from typing import List
+
+class CLIPScoreMetric(CLIPScore):
+    def update(self, images: Tensor | List[Tensor], text: str | list[str]) -> None:
+        """Update CLIP score on a batch of images and text.
+
+        Args:
+            images: Either a single [N, C, H, W] tensor or a list of [C, H, W] tensors, in the [-1, 1] range
+            text: Either a single caption or a list of captions
+
+        Raises:
+            ValueError:
+                If not all images have format [C, H, W]
+            ValueError:
+                If the number of images and captions do not match
+
+        """
+        images = (images+1)/2 # [-1,1] -> [0,1]
+        score, n_samples = _clip_score_update(images, text, self.model, self.processor)
+        self.score += score.sum(0)
+        self.n_samples += n_samples

hcpdiff/evaluate/previewer.py

@@ -6,32 +6,49 @@ from rainbowneko.utils import to_cuda
 
 from hcpdiff.models.wrapper import SD15Wrapper
 from accelerate.hooks import remove_hook_from_module
+from typing import Dict
+from types import ModuleType
 
 class HCPPreviewer(WorkflowPreviewer):
+    def __init__(self, parser, cfgs_raw, workflow: str | ModuleType | Dict, ds_name=None, interval=100, trainer=None,
+                 mixed_precision=None, seed=42, **cfgs):
+        super().__init__(parser, cfgs_raw, workflow, ds_name=ds_name, interval=interval, trainer=trainer,
+                         mixed_precision=mixed_precision, seed=seed, **cfgs)
+        if trainer is None:
+            self.pt_trainable = False
+        else:
+            self.emb_pt = trainer.cfgs.emb_pt
+            self.pt_trainable = trainer.pt_trainable
 
     @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:
+    def evaluate(self, step: int, prefix='eval/'):
+        if step%self.interval != 0 or not self.is_local_main_process:
             return
 
         # record training layers
-        training_layers = [layer for layer in model.modules() if layer.training]
+        if self.model_wrapper is not None:
+            training_layers = [layer for layer in self.model_raw.modules() if layer.training]
+            self.model_wrapper.eval()
+            model = self.model_raw
+        else:
+            training_layers = []
+            model = None
 
-        model.eval()
-        self.trainer.loggers.info(f'Preview')
+        if self.loggers is not None:
+            self.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 = Path(self.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)
+                                          device=self.device, dtype=self.weight_dtype, preview_root=preview_root, preview_step=step,
+                                          world_size=self.world_size, local_rank=self.local_rank,
+                                          emb_hook=self.emb_pt.embedding_hook if self.pt_trainable else None)
 
         # restore model states
         if model.vae is not None:
@@ -42,11 +59,11 @@ class HCPPreviewer(WorkflowPreviewer):
                 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:
+        if 'emb_hook' in states and not self.pt_trainable:
             states['emb_hook'].remove()
 
-        if self.trainer.pt_trainable:
-            self.trainer.cfgs.emb_pt.embedding_hook.N_repeats = N_repeats
+        if self.pt_trainable:
+            self.emb_pt.embedding_hook.N_repeats = N_repeats
 
         model.tokenizer.N_repeats = N_repeats
         model.text_enc_hook.N_repeats = N_repeats

hcpdiff/loss/base.py

@@ -6,36 +6,19 @@ 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')
+        self.target_type = getattr(loss, 'target_type', None)
 
-    def get_target(self, pred_type, model_pred, x_0, noise, x_t, sigma, noise_sampler, **kwargs):
+    def get_target(self, model_pred, x_0, noise, x_t, timesteps, 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}")
+        target = noise_sampler.get_target(x_0, x_t, timesteps, eps=noise, 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
+        # Convert pred_type for target_type
+        pred = noise_sampler.pred_for_target(model_pred, x_t, timesteps, eps=noise, target_type=self.target_type)
+        return 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_cvt, target = self.get_target(**pred)
+        pred['model_pred'] = pred_cvt
         pred['target'] = target
-        loss = super().forward(pred, inputs) * self.weight # [B,*,*,*]
+        loss = super().forward(pred, inputs) # [B,*,*,*]
         return loss.mean()

hcpdiff/loss/weighting.py

@@ -7,6 +7,11 @@ class LossWeight(nn.Module):
         super().__init__()
         self.loss = loss
 
+    def get_c_out(self, pred):
+        t = pred['timesteps']
+        noise_sampler = pred['noise_sampler']
+        return noise_sampler.sigma_scheduler.c_out(t)
+
     def get_weight(self, pred, inputs):
         '''
 
@@ -25,13 +30,19 @@ class LossWeight(nn.Module):
 
 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)
+        noise_sampler = pred['noise_sampler']
+        c_out = self.get_c_out(pred)
+        target_type = getattr(self.loss, 'target_type', None) or noise_sampler.target_type
+        if target_type == 'eps':
+            w_snr = 1
+        elif target_type == "x0":
+            w_snr = (1./c_out**2).float()
+        elif target_type == "velocity":
+            w_snr = (1./(1-c_out)**2).float()
         else:
-            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {target_type}")
+
+        return w_snr.view(-1, 1, 1, 1)
 
 class MinSNRWeight(LossWeight):
     def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
@@ -39,13 +50,18 @@ class MinSNRWeight(LossWeight):
         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()
+        noise_sampler = pred['noise_sampler']
+        c_out = self.get_c_out(pred)
+        target_type = getattr(self.loss, 'target_type', None) or noise_sampler.target_type
+        if target_type == 'eps':
+            w_snr = (self.gamma*c_out**2).clip(max=1).float()
+        elif target_type == "x0":
+            w_snr = (1./c_out**2).clip(max=self.gamma).float()
+        elif target_type == "velocity":
+            w_v = 1/(1-c_out)**2
+            w_snr = (self.gamma*c_out**2/w_v).clip(max=w_v).float()
         else:
-            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {target_type}")
 
         return w_snr.view(-1, 1, 1, 1)
 
@@ -55,12 +71,14 @@ class EDMWeight(LossWeight):
         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()
+        c_out = self.get_c_out(pred)
+        noise_sampler = pred['noise_sampler']
+        target_type = getattr(self.loss, 'target_type', None) or noise_sampler.target_type
+        if target_type == 'edm':
+            w_snr = 1
+        elif target_type == "x0":
+            w_snr = (1./c_out**2).float()
         else:
-            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {target_type}")
 
         return w_snr.view(-1, 1, 1, 1)

hcpdiff/models/lora_base_patch.py

@@ -34,6 +34,32 @@ class LoraPatchContainer(PatchPluginContainer):
 
         return self[name].post_forward(x, self._host.weight, weight_, self._host.bias, bias_)
 
+    @property
+    def weight(self):
+        weight_ = None
+        for name in self.plugin_names:
+            if weight_ is None:
+                weight_ = self[name].get_weight()
+            else:
+                weight_ = weight_+self[name].get_weight()
+        return self._host.weight + weight_
+
+    @property
+    def bias(self):
+        bias_ = None
+        for name in self.plugin_names:
+            if bias_ is None:
+                bias_ = self[name].get_bias()
+            else:
+                bias_ = bias_+self[name].get_bias()
+
+        if self._host.bias is not None:
+            if bias_ is None:
+                bias_ = self._host.bias
+            else:
+                bias_ = self._host.bias + bias_
+        return bias_
+
 class LoraBlock(PatchPluginBlock):
     container_cls = LoraPatchContainer
     wrapable_classes = (nn.Linear, nn.Conv2d)

hcpdiff/models/wrapper/sd.py

@@ -17,7 +17,7 @@ from ..cfg_context import CFGContext
 
 class SD15Wrapper(BaseWrapper):
     def __init__(self, denoiser: UNet2DConditionModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
-                 pred_type='eps', TE_hook_cfg:TEHookCFG=SD15_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
+                 TE_hook_cfg:TEHookCFG=SD15_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
         super().__init__()
         self.key_mapper_in = self.build_mapper(key_map_in, None, (
             'prompt -> prompt_ids', 'image -> image', 'attn_mask -> attn_mask', 'position_ids -> position_ids', 'neg_prompt -> neg_prompt_ids',
@@ -31,8 +31,6 @@ class SD15Wrapper(BaseWrapper):
         self.tokenizer = tokenizer
         self.min_attnmask = min_attnmask
 
-        self.pred_type = pred_type
-
         self.TE_hook_cfg = TEHookCFG.create(TE_hook_cfg)
         self.cfg_context = cfg_context
         self.tokenizer.N_repeats = self.TE_hook_cfg.tokenizer_repeats
@@ -93,8 +91,9 @@ class SD15Wrapper(BaseWrapper):
                       plugin_input={}, **kwargs):
         # input prepare
         x_0 = self.get_latents(image)
-        x_t, noise, sigma, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
-        x_t_in = x_t*self.noise_sampler.c_in(sigma).to(dtype=x_t.dtype)
+        x_t, noise, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
+        x_t_in = x_t*self.noise_sampler.sigma_scheduler.c_in(timesteps).to(dtype=x_t.dtype).view(-1,1,1,1)
+        t_in = self.noise_sampler.sigma_scheduler.c_noise(timesteps)
 
         if neg_prompt_ids:
             prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
@@ -104,15 +103,14 @@ class SD15Wrapper(BaseWrapper):
                 position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
 
         # model forward
-        x_t_in, timesteps = self.cfg_context.pre(x_t_in, timesteps)
-        encoder_hidden_states = self.forward_TE(prompt_ids, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+        x_t_in, t_in = self.cfg_context.pre(x_t_in, t_in)
+        encoder_hidden_states = self.forward_TE(prompt_ids, t_in, attn_mask=attn_mask, position_ids=position_ids,
                                                 plugin_input=plugin_input, **kwargs)
-        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, t_in, attn_mask=attn_mask, position_ids=position_ids,
                                            plugin_input=plugin_input, **kwargs)
         model_pred = self.cfg_context.post(model_pred)
 
-        return dict(model_pred=model_pred, noise=noise, sigma=sigma, timesteps=timesteps, x_0=x_0, x_t=x_t, pred_type=self.pred_type,
-                    noise_sampler=self.noise_sampler)
+        return dict(model_pred=model_pred, noise=noise, timesteps=timesteps, x_0=x_0, x_t=x_t, noise_sampler=self.noise_sampler)
 
     def forward(self, ds_name=None, **kwargs):
         model_args, model_kwargs = self.get_map_data(self.key_mapper_in, kwargs, ds_name)
@@ -156,8 +154,8 @@ class SD15Wrapper(BaseWrapper):
 
 class SDXLWrapper(SD15Wrapper):
     def __init__(self, denoiser: UNet2DConditionModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
-                 pred_type='eps', TE_hook_cfg:TEHookCFG=SDXL_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
-        super().__init__(denoiser, TE, vae, noise_sampler, tokenizer, min_attnmask, pred_type, TE_hook_cfg, cfg_context, key_map_in, key_map_out)
+                 TE_hook_cfg:TEHookCFG=SDXL_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
+        super().__init__(denoiser, TE, vae, noise_sampler, tokenizer, min_attnmask, TE_hook_cfg, cfg_context, key_map_in, key_map_out)
         self.key_mapper_in = self.build_mapper(key_map_in, None, (
             'prompt -> prompt_ids', 'image -> image', 'attn_mask -> attn_mask', 'position_ids -> position_ids', 'neg_prompt -> neg_prompt_ids',
             'neg_attn_mask -> neg_attn_mask', 'neg_position_ids -> neg_position_ids', 'plugin_input -> plugin_input', 'coord -> crop_info'))
@@ -195,8 +193,9 @@ class SDXLWrapper(SD15Wrapper):
                       crop_info=None, plugin_input={}):
         # input prepare
         x_0 = self.get_latents(image)
-        x_t, noise, sigma, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
-        x_t_in = x_t*self.noise_sampler.c_in(sigma).to(dtype=x_t.dtype)
+        x_t, noise, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
+        x_t_in = x_t*self.noise_sampler.sigma_scheduler.c_in(timesteps).to(dtype=x_t.dtype).view(-1,1,1,1)
+        t_in = self.noise_sampler.sigma_scheduler.c_noise(timesteps)
 
         if neg_prompt_ids:
             prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
@@ -206,13 +205,12 @@ class SDXLWrapper(SD15Wrapper):
                 position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
 
         # model forward
-        x_t_in, timesteps = self.cfg_context.pre(x_t_in, timesteps)
-        encoder_hidden_states, pooled_output = self.forward_TE(prompt_ids, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+        x_t_in, t_in = self.cfg_context.pre(x_t_in, t_in)
+        encoder_hidden_states, pooled_output = self.forward_TE(prompt_ids, t_in, attn_mask=attn_mask, position_ids=position_ids,
                                                                plugin_input=plugin_input)
         added_cond_kwargs = {"text_embeds":pooled_output[-1], "time_ids":crop_info}
-        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, timesteps, added_cond_kwargs=added_cond_kwargs,
+        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, t_in, added_cond_kwargs=added_cond_kwargs,
                                            attn_mask=attn_mask, position_ids=position_ids, plugin_input=plugin_input)
         model_pred = self.cfg_context.post(model_pred)
 
-        return dict(model_pred=model_pred, noise=noise, sigma=sigma, timesteps=timesteps, x_0=x_0, x_t=x_t, pred_type=self.pred_type,
-                    noise_sampler=self.noise_sampler)
+        return dict(model_pred=model_pred, noise=noise, timesteps=timesteps, x_0=x_0, x_t=x_t, noise_sampler=self.noise_sampler)

hcpdiff/trainer_ac.py

@@ -4,8 +4,8 @@ import warnings
 import torch
 from rainbowneko.parser import load_config_with_cli
 from rainbowneko.ckpt_manager import NekoSaver
-from rainbowneko.train import Trainer
-from rainbowneko.utils import xformers_available, is_dict
+from rainbowneko.train.trainer import Trainer
+from rainbowneko.utils import xformers_available, is_dict, weight_dtype_map
 from hcpdiff.ckpt_manager import EmbFormat
 
 class HCPTrainer(Trainer):
@@ -17,7 +17,7 @@ class HCPTrainer(Trainer):
                 warnings.warn("xformers is not available. Make sure it is installed correctly")
 
         if self.model_wrapper.vae is not None:
-            self.vae_dtype = self.weight_dtype_map.get(self.cfgs.model.get('vae_dtype', None), torch.float32)
+            self.vae_dtype = weight_dtype_map.get(self.cfgs.model.get('vae_dtype', None), torch.float32)
             self.model_wrapper.set_dtype(self.weight_dtype, self.vae_dtype)
 
         if self.cfgs.model.gradient_checkpointing:
@@ -44,10 +44,12 @@ class HCPTrainer(Trainer):
 
     def save_model(self, from_raw=False):
         NekoSaver.save_all(
-            self.model_raw,
-            plugin_groups={**self.all_plugin, 'embs': self.train_pts},
             cfg=self.ckpt_saver,
+            model=self.model_raw,
+            plugin_groups=self.all_plugin,
+            embs=self.train_pts,
             model_ema=getattr(self, "ema_model", None),
+            optimizer=self.optimizer,
             name_template=f'{{}}-{self.real_step}',
         )
 

hcpdiff/trainer_ac_single.py

@@ -1,12 +1,7 @@
 import argparse
-import sys
-from functools import partial
-
-import torch
-from accelerate import Accelerator
-from loguru import logger
 
 from rainbowneko.train.trainer import TrainerSingleCard
+
 from hcpdiff.trainer_ac import HCPTrainer, load_config_with_cli
 
 class HCPTrainerSingleCard(TrainerSingleCard, HCPTrainer):

hcpdiff/utils/__init__.py

@@ -1,2 +1,3 @@
 from .utils import *
-from .net_utils import *
+from .net_utils import *
+from .torch_utils import invert_func

hcpdiff/utils/torch_utils.py

@@ -0,0 +1,25 @@
+import torch
+
+def invert_func(func, y, x_min=0.0, x_max=1.0, tol=1e-5, max_iter=100):
+    """
+    y: [B]
+    :return: x [B]
+    """
+    y = y.to(dtype=torch.float32)
+    left = torch.full_like(y, x_min)
+    right = torch.full_like(y, x_max)
+
+    for _ in range(max_iter):
+        mid = (left+right)/2
+        val = func(mid)
+
+        too_large = val>y
+        too_small = ~too_large
+
+        left = torch.where(too_small, mid, left)
+        right = torch.where(too_large, mid, right)
+
+        if torch.all(torch.abs(val-y)<tol):
+            break
+
+    return (left+right)/2

hcpdiff/workflow/__init__.py

@@ -1,5 +1,5 @@
 from .diffusion import InputFeederAction, MakeLatentAction, DenoiseAction, SampleAction, DiffusionStepAction, \
-    X0PredAction, SeedAction, MakeTimestepsAction, PrepareDiffusionAction, time_iter
+    X0PredAction, SeedAction, MakeTimestepsAction, PrepareDiffusionAction, time_iter, DiffusionActions
 from .text import TextEncodeAction, TextHookAction, AttnMultTextEncodeAction
 from .vae import EncodeAction, DecodeAction
 from .io import BuildModelsAction, SaveImageAction, LoadImageAction