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

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.1.dist-info → hcpdiff-2.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hcpdiff
-Version: 2.2.1
+Version: 2.3.1
 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==1.6
+Requires-Dist: rainbowneko>=1.9
 Requires-Dist: diffusers
 Requires-Dist: matplotlib
 Requires-Dist: pyarrow

{hcpdiff-2.2.1.dist-info → hcpdiff-2.3.1.dist-info}/RECORD

@@ -1,16 +1,16 @@
 hcpdiff/__init__.py,sha256=dwNwrEgvG4g60fGMG6b50K3q3AWD1XCfzlIgbxkSUpE,177
 hcpdiff/train_colo.py,sha256=EsuNSzLBvGTZWU_LEk0JpP-F5eNW0lwkawIRAX38jmE,9250
-hcpdiff/trainer_ac.py,sha256=scH3FU0onCQtwLiy0-pcrhuowTZob3fLQqRP52iwY0c,2717
-hcpdiff/trainer_ac_single.py,sha256=0PIC5EScqcxp49EaeIWq4KS5K_09OZfKajqbFu-hUb8,1108
+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=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/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,41 +25,44 @@ 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=kKdESVqAxNlBhhz12PvwrpHJBea80OUFzDDMHwiulVs,6710
-hcpdiff/easy/cfg/sdxl_train.py,sha256=FUWE_hRJdQc9Qd9J6730jAyK0H4EIKS7-3BSufCItXU,4275
+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=O0XZqid01OrB0dHY7hCiBvdU2026SvZ38yfQaF2TWrs,8018
@@ -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.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-hcpdiff-2.2.1.dist-info/METADATA,sha256=f96Tc90K5WTBbJ35wWJw60G2JR46eGpUvQSaPIysVDg,10323
-hcpdiff-2.2.1.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-hcpdiff-2.2.1.dist-info/entry_points.txt,sha256=_4VRsEsEWOhHfzBDu9bx8Wh_S8Wi4ZTHpI0n6rU0J-I,258
-hcpdiff-2.2.1.dist-info/top_level.txt,sha256=shyf78x-HVgykYpsmY22mKG0xIc7Qk30fDMdavdYWQ8,8
-hcpdiff-2.2.1.dist-info/RECORD,,
+hcpdiff-2.3.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+hcpdiff-2.3.1.dist-info/METADATA,sha256=zaJHhKQiezDTvyv-IIoRHf4VCv0z2gU9fq0sVi9XhTg,10323
+hcpdiff-2.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+hcpdiff-2.3.1.dist-info/entry_points.txt,sha256=_4VRsEsEWOhHfzBDu9bx8Wh_S8Wi4ZTHpI0n6rU0J-I,258
+hcpdiff-2.3.1.dist-info/top_level.txt,sha256=shyf78x-HVgykYpsmY22mKG0xIc7Qk30fDMdavdYWQ8,8
+hcpdiff-2.3.1.dist-info/RECORD,,

{hcpdiff-2.2.1.dist-info → hcpdiff-2.3.1.dist-info}/WHEEL

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

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

hcpdiff/diffusion/sampler/ddpm.py

@@ -1,20 +0,0 @@
-import torch
-
-from .base import BaseSampler
-from .sigma_scheduler import SigmaScheduler
-
-class DDPMSampler(BaseSampler):
-    def __init__(self, sigma_scheduler: SigmaScheduler, generator: torch.Generator=None):
-        super().__init__(sigma_scheduler, generator)
-
-    def c_in(self, sigma):
-        return 1./(sigma**2+1).sqrt()
-
-    def c_out(self, sigma):
-        return -sigma
-
-    def c_skip(self, sigma):
-        return 1.
-
-    def denoise(self, x, sigma, eps=None, generator=None):
-        raise NotImplementedError

hcpdiff/diffusion/sampler/edm.py

@@ -1,22 +0,0 @@
-import torch
-
-from .base import BaseSampler
-from .sigma_scheduler import SigmaScheduler
-
-class EDMSampler(BaseSampler):
-    def __init__(self, sigma_scheduler: SigmaScheduler, generator: torch.Generator = None, sigma_data: float = 1.0, sigma_thr=1000):
-        super().__init__(sigma_scheduler, generator)
-        self.sigma_data = sigma_data
-        self.sigma_thr = sigma_thr
-
-    def c_in(self, sigma):
-        return 1/(sigma**2+self.sigma_data**2).sqrt()
-
-    def c_out(self, sigma):
-        return (sigma*self.sigma_data)/(sigma**2+self.sigma_data**2).sqrt()
-
-    def c_skip(self, sigma):
-        return self.sigma_data**2/(sigma**2+self.sigma_data**2)
-
-    def denoise(self, x, sigma, eps=None, generator=None):
-        raise NotImplementedError