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

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:
@@ -42,19 +42,14 @@ class HCPTrainer(Trainer):
     def pt_trainable(self):
         return self.cfgs.emb_pt is not None
 
-    def get_loss(self, ds_name, model_pred, inputs):
-        loss = super().get_loss(ds_name, model_pred, inputs)
-        # make DDP happy
-        if len(self.train_pts)>0:
-            loss = loss+0*sum([emb.mean() for emb in self.train_pts.values()])
-        return loss
-
     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/trainer_deepspeed.py

@@ -0,0 +1,47 @@
+import argparse
+import warnings
+
+import torch
+from rainbowneko.ckpt_manager import NekoPluginSaver
+from rainbowneko.train.trainer import TrainerDeepspeed
+from rainbowneko.utils import xformers_available
+
+from hcpdiff.trainer_ac import HCPTrainer, load_config_with_cli
+
+class HCPTrainerDeepspeed(TrainerDeepspeed, HCPTrainer):
+    def config_model(self):
+        if self.cfgs.model.enable_xformers:
+            if xformers_available:
+                self.model_wrapper.enable_xformers()
+            else:
+                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.model_wrapper.set_dtype(self.weight_dtype, self.vae_dtype)
+
+        if self.cfgs.model.gradient_checkpointing:
+            self.model_wrapper.enable_gradient_checkpointing()
+
+        if self.is_local_main_process:
+            for saver in self.ckpt_saver.values():
+                if isinstance(saver, NekoPluginSaver):
+                    saver.plugin_from_raw = True
+
+def hcp_train():
+    import subprocess
+    parser = argparse.ArgumentParser(description='HCP-Diffusion Launcher')
+    parser.add_argument('--launch_cfg', type=str, default='cfgs/launcher/deepspeed.yaml')
+    args, train_args = parser.parse_known_args()
+
+    subprocess.run(["accelerate", "launch", '--config_file', args.launch_cfg, "-m",
+                       "hcpdiff.trainer_deepspeed"]+train_args, check=True)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='HCP Diffusion Trainer for DeepSpeed')
+    parser.add_argument("--cfg", type=str, default=None, required=True)
+    args, cfg_args = parser.parse_known_args()
+
+    parser, conf = load_config_with_cli(args.cfg, args_list=cfg_args)  # skip --cfg
+    trainer = HCPTrainerDeepspeed(parser, conf)
+    trainer.train()

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

hcpdiff/workflow/diffusion.py

@@ -6,7 +6,7 @@ import torch
 from hcpdiff.diffusion.sampler import BaseSampler, DiffusersSampler
 from hcpdiff.utils import prepare_seed
 from hcpdiff.utils.net_utils import get_dtype, to_cuda
-from rainbowneko.infer import BasicAction
+from rainbowneko.infer import BasicAction, Actions
 from torch.cuda.amp import autocast
 
 try:
@@ -32,8 +32,9 @@ class SeedAction(BasicAction):
         self.seed = seed
         self.bs = bs
 
-    def forward(self, device, seed=None, **states):
-        bs = states['prompt_embeds'].shape[0]//2 if 'prompt_embeds' in states else self.bs
+    def forward(self, device, seed=None, bs=None, **states):
+        if bs is None:
+            bs = states['prompt_embeds'].shape[0]//2 if 'prompt_embeds' in states else self.bs
         seed = seed or self.seed
         if seed is None:
             seeds = [None]*bs
@@ -155,15 +156,16 @@ class DenoiseAction(BasicAction):
 
         with autocast(enabled=amp is not None, dtype=get_dtype(amp)):
             latent_model_input = torch.cat([latents]*2) if self.guidance_scale>1 else latents
-            latent_model_input = noise_sampler.c_in(t)*latent_model_input
+            latent_model_input = noise_sampler.sigma_scheduler.c_in(t)*latent_model_input
+            t_in = noise_sampler.sigma_scheduler.c_noise(t)
 
             if text_embeds is None:
-                noise_pred = denoiser(latent_model_input, t, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
+                noise_pred = denoiser(latent_model_input, t_in, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
                                   cross_attention_kwargs=cross_attention_kwargs, ).sample
             else:
                 added_cond_kwargs = {"text_embeds":text_embeds, "time_ids":crop_info}
                 # predict the noise residual
-                noise_pred = denoiser(latent_model_input, t, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
+                noise_pred = denoiser(latent_model_input, t_in, prompt_embeds, encoder_attention_mask=encoder_attention_mask,
                                   cross_attention_kwargs=cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs).sample
 
             # perform guidance
@@ -189,10 +191,28 @@ class DiffusionStepAction(BasicAction):
         states = self.act_noise_pred(denoiser=denoiser, noise_sampler=noise_sampler, **states)
         states = self.act_sample(**states)
         return states
+    
+class DiffusionActions(Actions):
+    def __init__(self, actions: List[BasicAction], clean_latent=True, seed_inc=True, key_map_in=None, key_map_out=None):
+        super().__init__(actions, key_map_in=key_map_in, key_map_out=key_map_out)
+        self.clean_latent = clean_latent
+        self.seed_inc = seed_inc
+
+    def forward(self, **states):
+        states = super().forward(**states)
+        if self.seed_inc and 'seed' in states:
+            bs = states['latents'].shape[0]
+            states['seed'] = states['seed'] + bs
+        if self.clean_latent:
+            states.pop('noise_pred', None)
+            states.pop('latents', None)
+            states.pop('prompt', None)
+            states.pop('negative_prompt', None)
+        return states
 
 class X0PredAction(BasicAction):
     def forward(self, latents, noise_sampler: BaseSampler, t, noise_pred, **states):
-        latents_x0 = noise_sampler.eps_to_x0(noise_pred, latents, t)
+        latents_x0 = noise_sampler.pred_for_target(noise_pred, latents, t, target_type='x0')
         return {'latents_x0':latents_x0}
 
 def time_iter(timesteps, **states):

hcpdiff/workflow/io.py

@@ -1,6 +1,7 @@
 import os
 from functools import partial
 from typing import List, Union
+from addict import Addict
 
 import torch
 from hcpdiff.utils import to_validate_file
@@ -9,6 +10,8 @@ from rainbowneko.ckpt_manager import NekoLoader
 from rainbowneko.infer import BasicAction
 from rainbowneko.infer import LoadImageAction as Neko_LoadImageAction
 from rainbowneko.utils.img_size_tool import types_support
+from rainbowneko import _share
+from rainbowneko.utils import is_dict
 
 class BuildModelsAction(BasicAction):
     def __init__(self, model_loader: partial[NekoLoader.load], dtype: str=torch.float32, device='cuda', key_map_in=None, key_map_out=None):
@@ -23,6 +26,14 @@ class BuildModelsAction(BasicAction):
         else:
             model = self.model_loader(dtype=self.dtype, device=self.device)
 
+            # Callback for TokenizerHandler
+            if is_dict(model):
+                model_wrapper = Addict(model)
+            else:
+                model_wrapper = model
+            for callback in _share.model_callbacks:
+                callback(model_wrapper)
+
         if isinstance(model, dict):
             return model
         else:
@@ -33,12 +44,13 @@ class LoadImageAction(Neko_LoadImageAction):
         super().__init__(image_paths, image_transforms, key_map_in, key_map_out)
 
 class SaveImageAction(BasicAction):
-    def __init__(self, save_root: str, image_type: str = 'png', quality: int = 95, save_cfg=True, key_map_in=None, key_map_out=None):
+    def __init__(self, save_root: str, image_type: str = 'png', quality: int = 95, save_cfg=True, save_txt=False, key_map_in=None, key_map_out=None):
         super().__init__(key_map_in, key_map_out)
         self.save_root = save_root
         self.image_type = image_type
         self.quality = quality
         self.save_cfg = save_cfg
+        self.save_txt = save_txt
 
         os.makedirs(save_root, exist_ok=True)
 
@@ -47,10 +59,15 @@ class SaveImageAction(BasicAction):
         num_img_exist = max([0]+[int(x.split('-', 1)[0]) for x in os.listdir(save_root) if x.rsplit('.', 1)[-1] in types_support])+1
 
         for bid, (p, pn, img) in enumerate(zip(prompt, negative_prompt, images)):
-            img_path = os.path.join(save_root, f"{preview_step or num_img_exist}-{seeds[bid]}-{to_validate_file(prompt[0])}.{self.image_type}")
+            img_path = os.path.join(save_root, f"{preview_step or num_img_exist}-{seeds[bid]}-{to_validate_file(p)}.{self.image_type}")
             img.save(img_path, quality=self.quality)
-            num_img_exist += 1
 
             if self.save_cfg:
                 cfgs.seed = seeds[bid]
                 parser.save_configs(cfgs, os.path.join(save_root, f"{preview_step or num_img_exist}-{seeds[bid]}-info"))
+
+            if self.save_txt:
+                txt_path = os.path.join(save_root, f"{preview_step or num_img_exist}-{seeds[bid]}-{to_validate_file(prompt[0])}.txt")
+                with open(txt_path, 'w') as f:
+                    f.write(p)
+            num_img_exist += 1

hcpdiff/workflow/text.py

@@ -38,7 +38,7 @@ class TextHookAction(BasicAction):
         return {'te_hook':te_hook, 'emb_hook':emb_hook, 'token_ex':token_ex}
 
 class TextEncodeAction(BasicAction):
-    def __init__(self, prompt: Union[List, str], negative_prompt: Union[List, str], bs: int = None, key_map_in=None, key_map_out=None):
+    def __init__(self, prompt: List|str|None, negative_prompt: List|str|None, bs: int = None, key_map_in=None, key_map_out=None):
         super().__init__(key_map_in, key_map_out)
         if isinstance(prompt, str) and bs is not None:
             prompt = [prompt]*bs
@@ -73,6 +73,11 @@ class AttnMultTextEncodeAction(TextEncodeAction):
         prompt = prompt or self.prompt
         negative_prompt = negative_prompt or self.negative_prompt
 
+        if isinstance(negative_prompt, str) and isinstance(prompt, (list, tuple)):
+            negative_prompt = [negative_prompt]*len(prompt)
+        if isinstance(prompt, str) and isinstance(negative_prompt, (list, tuple)):
+            prompt = [prompt]*len(negative_prompt)
+
         if model_offload:
             to_cuda(TE)
 

{hcpdiff-2.2.dist-info → hcpdiff-2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hcpdiff
-Version: 2.2
+Version: 2.3
 Summary: A universal Diffusion toolbox
 Home-page: https://github.com/IrisRainbowNeko/HCP-Diffusion
 Author: Ziyi Dong
@@ -17,7 +17,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: rainbowneko
+Requires-Dist: rainbowneko>=1.9
 Requires-Dist: diffusers
 Requires-Dist: matplotlib
 Requires-Dist: pyarrow
@@ -262,9 +262,13 @@ hcp_run --cfg cfgs/workflow/text2img_cli.py \
     seed=42
 ```
 
-### Tutorials
+### 📚 Tutorials
 
-🚧 In Development
++ 🧠 [Model Training Guide](https://hcpdiff.readthedocs.io/en/latest/user_guides/train.html)
++ 🔧 [LoRA Training Tutorial](https://hcpdiff.readthedocs.io/enlatest/tutorial/lora.html)
++ 🎨 [Image Generation Guide](https://hcpdiff.readthedocs.io/en/latest/user_guides/workflow.html)
++ ⚙️ [Configuration File Explanation](https://hcpdiff.readthedocs.io/en/latest/user_guides/cfg.html)
++ 🧩 [Model Format Explanation](https://hcpdiff.readthedocs.io/en/latest/user_guides/model_format.html)
 
 ---
 

{hcpdiff-2.2.dist-info → hcpdiff-2.3.dist-info}/RECORD

@@ -1,16 +1,16 @@
 hcpdiff/__init__.py,sha256=dwNwrEgvG4g60fGMG6b50K3q3AWD1XCfzlIgbxkSUpE,177
 hcpdiff/train_colo.py,sha256=EsuNSzLBvGTZWU_LEk0JpP-F5eNW0lwkawIRAX38jmE,9250
-hcpdiff/train_deepspeed.py,sha256=PwyNukWi0of6TXy_VRDgBQSMLCZBhipO5g3Lq0nCYNk,2988
-hcpdiff/trainer_ac.py,sha256=6KAzo54in7ZRHud_rHjJdwRRZ4uWtc0B4SxVCxgcrmM,2990
-hcpdiff/trainer_ac_single.py,sha256=0PIC5EScqcxp49EaeIWq4KS5K_09OZfKajqbFu-hUb8,1108
-hcpdiff/ckpt_manager/__init__.py,sha256=Mn_5KOC4xbf2GcN6OXg_XdbF5wO9zWeER_1ZO_prKAI,256
-hcpdiff/ckpt_manager/ckpt.py,sha256=Pa3uXQbCi2T99mpV5fYddQ-OGHcpk8r1ll-0lmP_WXk,965
-hcpdiff/ckpt_manager/loader.py,sha256=Ch1xsZmseq4nyPhpox9-nebN-dZB4k0rqBEHos-ZLso,3245
+hcpdiff/trainer_ac.py,sha256=-owV-3_bvPxuQsZS2WaajBDh58HpftRtnx0GJkswqaY,2787
+hcpdiff/trainer_ac_single.py,sha256=zyZVrutLUbIJYW1HzUnQ_RnmIcDhbC7M_CT833PJH5w,993
+hcpdiff/trainer_deepspeed.py,sha256=7lGsiAstWuIlmhRMwWTcJCkoxzUaakVxBngKDnJdSJk,1947
+hcpdiff/ckpt_manager/__init__.py,sha256=r_sgjZWCLtdJrRkqqU6aPdfubXSYfPh2Z_Vf_XpZXXs,240
+hcpdiff/ckpt_manager/ckpt.py,sha256=2A093lT03M1ZsJIMWl376V165eh0TZwOgiGrz3LM73Q,1248
+hcpdiff/ckpt_manager/loader.py,sha256=6iZDUj-Vfc5T9eGdWfFMQw4n1GqyLqaLBolgAtgqPq8,3640
 hcpdiff/ckpt_manager/format/__init__.py,sha256=a3cdKkOTDgdVbDQwSC4mlxOigjX2hBvRb5_X7E3TQWs,237
-hcpdiff/ckpt_manager/format/diffusers.py,sha256=T81WN95Nj1il9DfQp9iioVn0uqFEWOlmdIYs2beNOFU,3769
+hcpdiff/ckpt_manager/format/diffusers.py,sha256=qhGbrKAaeLyjFzY-Lj4sL1THHFNrta41JGGMoXT-bCE,3761
 hcpdiff/ckpt_manager/format/emb.py,sha256=FrqfTfJ8H7f0Zw17NTWCP2AJtpsJI5oXR5IAd4NekhU,680
 hcpdiff/ckpt_manager/format/lora_webui.py,sha256=4y_T9RdmFTxWzsXd8guNjCiukmyILa5j4MPrhVIL4Qk,10017
-hcpdiff/ckpt_manager/format/sd_single.py,sha256=LpCAL_7nAVooCHTFznVVsNMku1G3C77NBORxxr8GDtQ,2328
+hcpdiff/ckpt_manager/format/sd_single.py,sha256=4DZLAl1RNC_nPxuW-lmrBlIMFUhpSTa7HGHgu7Yx8qk,2322
 hcpdiff/data/__init__.py,sha256=ZFKtanOoMo3G3eKUJPhysnHXnr8BNARERkcMB6B897U,292
 hcpdiff/data/dataset.py,sha256=1k4GldW13eVyqK_9hrQniqr3_XYAapnWF7iXl_1GXGg,877
 hcpdiff/data/cache/__init__.py,sha256=ToCmokYH6DghlSwm7HJFirPRIWJ0LkgzqVOYlgoAkQw,25
@@ -25,44 +25,47 @@ hcpdiff/data/source/text.py,sha256=VgI5Ouq986Yy1jwD2fZ9iBlsRciPCeARZmOPEZIcaQY,1
 hcpdiff/data/source/text2img.py,sha256=acYdolQhZUEpkd7tUAdNkCTVnPc1SMJOVTmGqFt9ZpE,1813
 hcpdiff/data/source/text2img_cond.py,sha256=yj1KpARA2rkjENutnnzC4uDkcU2Rye21FL2VdC25Hac,585
 hcpdiff/diffusion/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hcpdiff/diffusion/noise/__init__.py,sha256=seBpOtd0YsU53PqMn7Nyl_RtwoC-ONEIOX7v2XLGpZQ,93
+hcpdiff/diffusion/noise/__init__.py,sha256=D83EZ6bnc6Ucu4AZwE6rpmXCtwYfHHumeVq97brbnIE,47
 hcpdiff/diffusion/noise/pyramid_noise.py,sha256=KbpyMT1BHNIaAa7g5eECDkTttOMoMWVFmbP-ekBsuEY,1693
-hcpdiff/diffusion/noise/zero_terminal.py,sha256=EfVOaqrTCfw11AolDBl0LIOey3uQT1bDw2XKr2Bm434,1532
-hcpdiff/diffusion/sampler/__init__.py,sha256=pSHsKpLjscY5yLbdzHeBUeK9nFDuVeMIIeA_k6FQFdY,158
-hcpdiff/diffusion/sampler/base.py,sha256=2AuPVT2ZSXYt2etZmHMyNKuGlT5zn6KIkoMz4m5PGcs,2577
-hcpdiff/diffusion/sampler/ddpm.py,sha256=raqSuKsEPN1AEqRVCuBdMAOnKDoeJTRO17wtLBNJCf4,523
-hcpdiff/diffusion/sampler/diffusers.py,sha256=XIu-oIlT4LnAYI2-yyIoNeIMeSQe5YpeEvn-FkGVFnE,2684
-hcpdiff/diffusion/sampler/edm.py,sha256=5W4pv8hxQsPpJGiFBgElZxR3C9S8kWAhzGKejEVwq3I,753
-hcpdiff/diffusion/sampler/sigma_scheduler/__init__.py,sha256=kmIoWgsWqij6b7KYon3UOCSC07sRJCo-DPR6qkJwUd0,184
-hcpdiff/diffusion/sampler/sigma_scheduler/base.py,sha256=9-JI-jwf7xZoQUtrU0qfbjkhNZT8a_tmapLtwVbFUx0,381
-hcpdiff/diffusion/sampler/sigma_scheduler/ddpm.py,sha256=2PMIpg2K6CVoxew1y1pIqvCHbdggC-m3amrOYk15OdQ,8107
-hcpdiff/diffusion/sampler/sigma_scheduler/edm.py,sha256=fOPB3lgnS9uVo4oW26Fur_nc8X_wQ6mmUcbkKhnoQjs,1900
+hcpdiff/diffusion/sampler/VP.py,sha256=r0Q_RROEIeNNw93XrOD5htW78rfuoSxy1WBQEoQL83s,958
+hcpdiff/diffusion/sampler/__init__.py,sha256=Lrwg1us8qo943T7mdIXFDRXfKvnLhrzwmi6DrIKIiUA,135
+hcpdiff/diffusion/sampler/base.py,sha256=UbE_AmtvLg-Hr2bkYz8PvNWB63tvtacUvCIDm_W6opA,5484
+hcpdiff/diffusion/sampler/diffusers.py,sha256=wIMs8n3kdci1On0FUCV0si324ZE9zeRw_CxaHP8rdcs,2586
+hcpdiff/diffusion/sampler/sigma_scheduler/__init__.py,sha256=eiSmMBkXI_LfxnNrXj5XptcF0dGcPas--vWvqhFGlv8,273
+hcpdiff/diffusion/sampler/sigma_scheduler/base.py,sha256=UT4tbjFf80KYfU08y0hJf8h_Cl80a5MUhK5FsKLsqbY,2521
+hcpdiff/diffusion/sampler/sigma_scheduler/ddpm.py,sha256=SA8lXT6lucJot_rpJ84Wz-_uc5dXfb2QPoQgJHSKOj4,12999
+hcpdiff/diffusion/sampler/sigma_scheduler/edm.py,sha256=m1YlIyn61zfDjxLMcHvWs0nzULbHgXeB7WGKmTiaGSU,4127
+hcpdiff/diffusion/sampler/sigma_scheduler/flow.py,sha256=FtWpesUtSmFuiIGkrrVhYJweB7INZiw0atC64tc0Nk4,2020
+hcpdiff/diffusion/sampler/sigma_scheduler/zero_terminal.py,sha256=CCqQLkGo4omkxzFovYdZQzdZVwIxK3PiOitZFww8MHs,859
 hcpdiff/easy/__init__.py,sha256=-emoyCOZlLCu3KNMI8L4qapUEtEYFSoiGU6-rKv1at4,149
 hcpdiff/easy/sampler.py,sha256=dQSBkeGh71O0DAmZLhTHTbk1bY7XzyUCeW1oJO14A4I,1250
 hcpdiff/easy/cfg/__init__.py,sha256=SxHMWG6T2CXhX3dP0xizSMd9vFWPaZQDc4Gj4CF__yQ,253
-hcpdiff/easy/cfg/sd15_train.py,sha256=KtplqN-OhzdZjsX2s60J3XR6o7tRJ-QDx7Eqza_eDkM,6704
-hcpdiff/easy/cfg/sdxl_train.py,sha256=ZKfJ19IvR2dZqDNXULmhZEmqjE7qV4QYxSTvEhI7efQ,4269
+hcpdiff/easy/cfg/sd15_train.py,sha256=NtgsQLg1sd5JFmHU4nqMPOrvP7zmwo2x0MCspjVNQEY,7000
+hcpdiff/easy/cfg/sdxl_train.py,sha256=rVLLKVMKB_PHuum3dKQcBqL0uR8QhzmdRllM-pYnbK4,4534
 hcpdiff/easy/cfg/t2i.py,sha256=SnjFjZAKd9orjJr3RW5_N2_EIlW2Ree7JMvdNUAR9gc,9507
 hcpdiff/easy/model/__init__.py,sha256=CA-7r3R2Jgweekk1XNByFYttLolbWyUV2bCnXygcD8w,133
 hcpdiff/easy/model/cnet.py,sha256=m0NTH9V1kLzb5GybwBrSNT0KvTcRpPfGkzUeMz9jZZQ,1084
 hcpdiff/easy/model/loader.py,sha256=Tdx-lhQEYf2NYjVM1A5B8x6ZZpJKcXUkFIPIbr7h7XM,3456
-hcpdiff/evaluate/__init__.py,sha256=CtNzi8xdUWZBDBcP5TZTDMcRyykaOJhBIxTJgVuMabo,35
-hcpdiff/evaluate/previewer.py,sha256=QiYYiEJBKP06uL3wKLhnpIGUZuAkr9BuHxTE29obpXI,2432
+hcpdiff/evaluate/__init__.py,sha256=qWxV0D8Ho5uBj2YbaC_QFDnT49PSKPfh44m4ivkNbMM,108
+hcpdiff/evaluate/evaluator.py,sha256=9BZQBeC-N7p-ICx6Giw9v-2Tb9volMTDmeDfhj0nXJ0,2940
+hcpdiff/evaluate/previewer.py,sha256=-vE0YXVfos70CQMo9ZInw7xu3d88DlTfVLs4BzzkxfM,3140
+hcpdiff/evaluate/metrics/__init__.py,sha256=vE0nSvBtDBu9SomANvWcm2UHX56PhCYwhgrcmm_mKyo,39
+hcpdiff/evaluate/metrics/clip_score.py,sha256=rQgweu5QcqW3fPI3EXcNbrH2QCcSAekE3lpYk45P2M4,900
 hcpdiff/loss/__init__.py,sha256=2dwPczSiv3rB5fzOeYbl5ZHpMU-qXOQlXeOiXdxcxwM,173
-hcpdiff/loss/base.py,sha256=3bvgMbwyPOEA9iSkv0hRHw4VnKjkUCZAENNnDMFilYM,1780
+hcpdiff/loss/base.py,sha256=Vvpm-KZGH4n-gYIlnVAtPl1B799c7v0dJXJ5BBh3yO0,1112
 hcpdiff/loss/gw.py,sha256=0yi1kozuII3xZA6FnjOhINtvScWt1MyBZLBtMKmgojM,1224
 hcpdiff/loss/ssim.py,sha256=YofadvBkc6sklxBUx1p3ADw5OHOZPK3kaHz8FH5a6m4,1281
 hcpdiff/loss/vlb.py,sha256=s78iBnXUiDWfGf7mYmhUnHqxqea5gSByKOoqBrX6bzU,3222
-hcpdiff/loss/weighting.py,sha256=UR2PyZ1JTNOydXMw4e1Fh52XmtwKaviPvddcmVKCTlI,2242
+hcpdiff/loss/weighting.py,sha256=9qzMnvCb6b5qx0p08GDSlkxmYEqQcNt79XdRBvfHmiI,2914
 hcpdiff/models/__init__.py,sha256=eQS7DPiGLiE1MFRkZj_17IY3IsfDUVcYpcOmhHb5B9o,472
 hcpdiff/models/cfg_context.py,sha256=e2B3K1KwJhzbD6xdJUOyNtl_XgQ0296XI3FHw3gvZF4,1502
 hcpdiff/models/container.py,sha256=z3p5TmQhxdzXSIfofz55_bmEhSsgUJsy1o9EcDs8Oeo,696
 hcpdiff/models/controlnet.py,sha256=VIkUzJCVpCqqQOtRSLQPfbcDy9CsXutxLeZB6PdZfA0,7809
 hcpdiff/models/lora_base.py,sha256=LGwBD9KP6qf4pgTx24i5-JLo4rDBQ6jFfterQKBjTbE,6758
-hcpdiff/models/lora_base_patch.py,sha256=WW3CULnROTxKXyynJiqirhHYCKN5JtxLhVpT5b7AUQg,6532
+hcpdiff/models/lora_base_patch.py,sha256=Tdb_b3TN_K-04nlUvcfBh6flPcbL9M4iP7jOVyb1jXQ,7271
 hcpdiff/models/lora_layers.py,sha256=O9W_Ue71lHj7Y_GbpioF4Hc3h2-z_zOqck93VYUra6s,7777
 hcpdiff/models/lora_layers_patch.py,sha256=GYFYsJD2VSLZfdnLma9CmQEHz09HROFJcc4wc_gs9f0,8198
-hcpdiff/models/text_emb_ex.py,sha256=a5QImxzvj0zWR12qXOPP9kmpESl8J9VLabA0W9D_i_c,7867
+hcpdiff/models/text_emb_ex.py,sha256=O0XZqid01OrB0dHY7hCiBvdU2026SvZ38yfQaF2TWrs,8018
 hcpdiff/models/textencoder_ex.py,sha256=JrTQ30Avx8tPbdr-Q6K5BvEWCEdsu8Z7eSOzMqpUuzg,8270
 hcpdiff/models/tokenizer_ex.py,sha256=zKUn4BY7b3yXwK9PWkZtQKJPyKYwUc07E-hwB9NQybs,2446
 hcpdiff/models/compose/__init__.py,sha256=lTNFTGg5csqvUuys22RqgjmWlk_7Okw6ZTsnTi1pqCg,217
@@ -72,7 +75,7 @@ hcpdiff/models/compose/compose_tokenizer.py,sha256=g3l0pOFv6p7Iigxm6Pqt_iTUXBlO1
 hcpdiff/models/compose/sdxl_composer.py,sha256=NtMGaFGZTfKsPJSVi2yT-UM6K1WKWtk99XxVmTcKlk8,2164
 hcpdiff/models/wrapper/__init__.py,sha256=HbGQmFnfccr-dtvZKjEv-pmR4cCnF4fwGLKS3tuG_OY,135
 hcpdiff/models/wrapper/pixart.py,sha256=nRUvHSHn4TYg_smC0xpeW-GtUgXss-MuaVPTHpMozDE,1147
-hcpdiff/models/wrapper/sd.py,sha256=D7VDI4OmbLTk9mzYta-C4LJjWfZmuBiDub4t8v1-M9o,11711
+hcpdiff/models/wrapper/sd.py,sha256=EywmVU2QzR74M_4eH_uXVW8HJNauyjwcZPU7rRAQ7eI,11666
 hcpdiff/models/wrapper/utils.py,sha256=NyebMoAPnrgcTHbiIocSD-eGdGdD-V1G_TQuWsRWufw,665
 hcpdiff/parser/__init__.py,sha256=-2dDZ2Ii4zoGQqDTme94q4PpJbBiV6HS5BsDASz4Xbo,33
 hcpdiff/parser/embpt.py,sha256=LgwZ0f0tLn3DrTo5ZpSCsZcA5330UpiW_sK96yEPmOM,1307
@@ -89,27 +92,28 @@ hcpdiff/tools/init_proj.py,sha256=XrXxxhIaItywG7HsrloJo-x8w9suZiY35daelzZvjrg,19
 hcpdiff/tools/lora_convert.py,sha256=So14WvSVIm6rU4m1XCajFXDnhq7abpZS95SLbaoyBFU,10058
 hcpdiff/tools/save_model.py,sha256=gbfYi_EfEBZEUcDjle6MDHA19sQWY0zA8_y_LMzHQ7M,428
 hcpdiff/tools/sd2diffusers.py,sha256=vB6OnBLw60sJkdpVZcYEPtKAZW1h8ErbSGSRq0uAiIk,16855
-hcpdiff/utils/__init__.py,sha256=VOLhdNq2mRyqmWxrssIWSZtR_PQ8rFwo2u0uq6GbLHA,45
+hcpdiff/utils/__init__.py,sha256=28K9Ui0uur-vHuUdlSyIBYijgu2b7rGOPXN2ogJu1z8,82
 hcpdiff/utils/colo_utils.py,sha256=JyLUvVnISa48CnryNLrgVxMo-jxu2UhBq70eYPrkjuI,837
 hcpdiff/utils/inpaint_pipe.py,sha256=CRy1MUlPmHifCAbZnKOP0qbLp2grn7ZbVeaB2qIA4ig,42862
 hcpdiff/utils/net_utils.py,sha256=gdwLYDNKV2t3SP0jBIO3d0HtY6E7jRaf_rmPT8gKZZE,9762
 hcpdiff/utils/pipe_hook.py,sha256=-UDX3FtZGl-bxSk13gdbPXc1OvtbCcpk_fvKxLQo3Ag,31987
+hcpdiff/utils/torch_utils.py,sha256=gBZCcDKZc0NGDQx6QeHuQePoZ82kQRhaL7oEdZIYGvU,573
 hcpdiff/utils/utils.py,sha256=hZnZP1IETgVpScxES0yIuRfc34TnzvAqmgOTK_56ssw,4976
-hcpdiff/workflow/__init__.py,sha256=t7Zyc0XFORdNvcwHp9AsCtEkhJ3l7Hm41ugngIL0Sag,867
-hcpdiff/workflow/diffusion.py,sha256=yzhqKA3019OPu1RKggrLoytMgm919qf6j9S85PYOwjQ,8644
+hcpdiff/workflow/__init__.py,sha256=i5s7QXo6wK9607KL0KTW4suE1c-HGJ5_EgnCdVLl3WM,885
+hcpdiff/workflow/diffusion.py,sha256=hKefBrVP6-025MhdrKOQMUhHxLaGqjpUKhR6WahYwh0,9549
 hcpdiff/workflow/fast.py,sha256=kZt7bKrvpFInSn7GzbkTkpoCSM0Z6IbDjgaDvcbFYf8,1024
 hcpdiff/workflow/flow.py,sha256=FFbFFOAXT4c31L5bHBEB_qeVGuBQDLYhq8kTD1chGNo,2548
-hcpdiff/workflow/io.py,sha256=aTrMR3s44apVJpnSyvZIabW2Op0tslk_Z9JFJl5svm0,2635
+hcpdiff/workflow/io.py,sha256=4oiE_PS3sOVYT8M6PDwvT5h9XzoKDMQR0n_4-Ktttys,3284
 hcpdiff/workflow/model.py,sha256=1gj5yOTefYTnGXVR6JPAfxIwuB69YwN6E-BontRcuyQ,2913
-hcpdiff/workflow/text.py,sha256=Z__SJHZyuaKyzkYJ6rbiAzOGRiYcCjwCGeqfpP1Jo7o,4336
+hcpdiff/workflow/text.py,sha256=XQvN4zzK7VaGxy4FDgSDeWh2jjk7UZU24moeRKAWXRE,4608
 hcpdiff/workflow/utils.py,sha256=xojaMG4lHsymslc8df5uiVXmmBVWpn_Phqka8qzJEWw,2226
 hcpdiff/workflow/vae.py,sha256=cingDPkIOc4qGpOwwhXJK4EQbGoIxO583pm6gGov5t8,3118
 hcpdiff/workflow/daam/__init__.py,sha256=ySIDaxloN-D3qM7OuVaG1BR3D-CibDoXYpoTgw0zUhU,59
 hcpdiff/workflow/daam/act.py,sha256=tHbsFWTYYU4bvcZOo1Bpi_z6ofpJatRYccl4vvf8wIA,2756
 hcpdiff/workflow/daam/hook.py,sha256=z9f9mBjKW21xuUZ-iQxQ0HbWOBXtZrisFB0VNMq6d0U,4383
-hcpdiff-2.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-hcpdiff-2.2.dist-info/METADATA,sha256=u52mZtA0hI2P_fObmJZRUkZZfnKFYg5c24f4p0trH0o,9833
-hcpdiff-2.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-hcpdiff-2.2.dist-info/entry_points.txt,sha256=86wPOMzsfWWflTJ-sQPLc7WG5Vtu0kGYBH9C_vR3ur8,207
-hcpdiff-2.2.dist-info/top_level.txt,sha256=shyf78x-HVgykYpsmY22mKG0xIc7Qk30fDMdavdYWQ8,8
-hcpdiff-2.2.dist-info/RECORD,,
+hcpdiff-2.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+hcpdiff-2.3.dist-info/METADATA,sha256=1aqZH8IwB7WjDBot_fADwyfLDVNiRZuZdub2-zezFck,10321
+hcpdiff-2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+hcpdiff-2.3.dist-info/entry_points.txt,sha256=_4VRsEsEWOhHfzBDu9bx8Wh_S8Wi4ZTHpI0n6rU0J-I,258
+hcpdiff-2.3.dist-info/top_level.txt,sha256=shyf78x-HVgykYpsmY22mKG0xIc7Qk30fDMdavdYWQ8,8
+hcpdiff-2.3.dist-info/RECORD,,

{hcpdiff-2.2.dist-info → hcpdiff-2.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

{hcpdiff-2.2.dist-info → hcpdiff-2.3.dist-info}/entry_points.txt

@@ -2,4 +2,5 @@
 hcp_run = rainbowneko.infer.infer_workflow:run_workflow
 hcp_train = hcpdiff.trainer_ac:hcp_train
 hcp_train_1gpu = hcpdiff.trainer_ac_single:hcp_train
+hcp_train_ds = hcpdiff.trainer_deepspeed:hcp_train
 hcpinit = hcpdiff.tools.init_proj:main

hcpdiff/diffusion/noise/zero_terminal.py

@@ -1,39 +0,0 @@
-import torch
-from ..sampler.sigma_scheduler import DDPMDiscreteSigmaScheduler
-
-class ZeroTerminalSampler:
-
-    @classmethod
-    def patch(cls, base_sampler):
-        assert isinstance(base_sampler.sigma_scheduler, DDPMDiscreteSigmaScheduler), "ZeroTerminalScheduler only works with DDPM SigmaScheduler"
-
-        alphas_cumprod = base_sampler.sigma_scheduler.alphas_cumprod
-        base_sampler.sigma_scheduler.alphas_cumprod = cls.rescale_zero_terminal_snr(alphas_cumprod)
-        base_sampler.sigma_scheduler.sigmas = ((1-alphas_cumprod)/alphas_cumprod).sqrt()
-
-
-    @staticmethod
-    def rescale_zero_terminal_snr(alphas_cumprod, thr=1e-4):
-        """
-        Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
-        Args:
-            alphas_cumprod (`torch.FloatTensor`)
-        Returns:
-            `torch.FloatTensor`: rescaled betas with zero terminal SNR
-        """
-        alphas_bar_sqrt = alphas_cumprod.sqrt()
-
-        # Store old values.
-        alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
-        alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
-
-        # Shift so the last timestep is zero.
-        alphas_bar_sqrt -= alphas_bar_sqrt_T
-
-        # Scale so the first timestep is back to the old value.
-        alphas_bar_sqrt *= alphas_bar_sqrt_0/(alphas_bar_sqrt_0-alphas_bar_sqrt_T)
-        alphas_bar_sqrt[-1] = thr # avoid nan sigma
-
-        # Convert alphas_bar_sqrt to betas
-        alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
-        return alphas_bar