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

hcpdiff/easy/cfg/sd15_train.py

@@ -47,7 +47,7 @@ def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
 
             optimizer=optimizer,
 
-            scheduler=ConstantLR(
+            lr_scheduler=ConstantLR(
                 _partial_=True,
                 warmup_steps=warmup_steps,
             ),
@@ -132,7 +132,7 @@ def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
 
             optimizer=optimizer,
 
-            scheduler=ConstantLR(
+            lr_scheduler=ConstantLR(
                 _partial_=True,
                 warmup_steps=warmup_steps,
             ),

hcpdiff/easy/cfg/sdxl_train.py

@@ -44,7 +44,7 @@ def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
 
             optimizer=optimizer,
 
-            scheduler=ConstantLR(
+            lr_scheduler=ConstantLR(
                 _partial_=True,
                 warmup_steps=warmup_steps,
             ),
@@ -128,7 +128,7 @@ def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
 
             optimizer=optimizer,
 
-            scheduler=ConstantLR(
+            lr_scheduler=ConstantLR(
                 _partial_=True,
                 warmup_steps=warmup_steps,
             ),

hcpdiff/models/text_emb_ex.py

@@ -126,6 +126,10 @@ class EmbeddingPTInterpHook(SinglePluginBlock):
         BOS = repeat(inputs_embeds[0,0,:], 'e -> r 1 e', r=self.N_repeats)
         EOS = repeat(inputs_embeds[0,-1,:], 'e -> r 1 e', r=self.N_repeats)
 
+        # make DDP happy
+        if len(self.emb_train) > 0:
+            BOS = BOS + sum(emb.mean()*0 for emb in self.emb_train if emb.requires_grad)
+
         replaced_embeds = []
         for item, rep_idxs, ids_raw in zip(inputs_embeds, rep_idxs_B, self.input_ids):
             # insert pt to embeddings

hcpdiff/trainer_ac.py

@@ -42,13 +42,6 @@ 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,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hcpdiff
-Version: 2.2
+Version: 2.2.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
+Requires-Dist: rainbowneko==1.6
 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.2.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 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.py,sha256=scH3FU0onCQtwLiy0-pcrhuowTZob3fLQqRP52iwY0c,2717
 hcpdiff/trainer_ac_single.py,sha256=0PIC5EScqcxp49EaeIWq4KS5K_09OZfKajqbFu-hUb8,1108
+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
@@ -40,8 +40,8 @@ hcpdiff/diffusion/sampler/sigma_scheduler/edm.py,sha256=fOPB3lgnS9uVo4oW26Fur_nc
 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=kKdESVqAxNlBhhz12PvwrpHJBea80OUFzDDMHwiulVs,6710
+hcpdiff/easy/cfg/sdxl_train.py,sha256=FUWE_hRJdQc9Qd9J6730jAyK0H4EIKS7-3BSufCItXU,4275
 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
@@ -62,7 +62,7 @@ hcpdiff/models/lora_base.py,sha256=LGwBD9KP6qf4pgTx24i5-JLo4rDBQ6jFfterQKBjTbE,6
 hcpdiff/models/lora_base_patch.py,sha256=WW3CULnROTxKXyynJiqirhHYCKN5JtxLhVpT5b7AUQg,6532
 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
@@ -107,9 +107,9 @@ 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.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.2.dist-info → hcpdiff-2.2.1.dist-info}/WHEEL

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

{hcpdiff-2.2.dist-info → hcpdiff-2.2.1.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/train_deepspeed.py

@@ -1,69 +0,0 @@
-import argparse
-import os
-import sys
-import warnings
-from functools import partial
-
-import torch
-
-from hcpdiff.ckpt_manager import CkptManagerPKL, CkptManagerSafe
-from hcpdiff.train_ac_old import Trainer, load_config_with_cli
-from hcpdiff.utils.net_utils import get_scheduler
-
-class TrainerDeepSpeed(Trainer):
-
-    def build_ckpt_manager(self):
-        self.ckpt_manager = self.ckpt_manager_map[self.cfgs.ckpt_type](plugin_from_raw=True)
-        if self.is_local_main_process:
-            self.ckpt_manager.set_save_dir(os.path.join(self.exp_dir, 'ckpts'), emb_dir=self.cfgs.tokenizer_pt.emb_dir)
-
-    @property
-    def unet_raw(self):
-        return self.accelerator.unwrap_model(self.TE_unet).unet if self.train_TE else self.accelerator.unwrap_model(self.TE_unet.unet)
-
-    @property
-    def TE_raw(self):
-        return self.accelerator.unwrap_model(self.TE_unet).TE if self.train_TE else self.TE_unet.TE
-
-    def get_loss(self, model_pred, target, timesteps, att_mask):
-        if att_mask is None:
-            att_mask = 1.0
-        if getattr(self.criterion, 'need_timesteps', False):
-            loss = (self.criterion(model_pred.float(), target.float(), timesteps)*att_mask).mean()
-        else:
-            loss = (self.criterion(model_pred.float(), target.float())*att_mask).mean()
-        return loss
-
-    def build_optimizer_scheduler(self):
-        # set optimizer
-        parameters, parameters_pt = self.get_param_group_train()
-
-        if len(parameters_pt)>0:  # do prompt-tuning
-            cfg_opt_pt = self.cfgs.train.optimizer_pt
-            # if self.cfgs.train.scale_lr_pt:
-            #     self.scale_lr(parameters_pt)
-            assert isinstance(cfg_opt_pt, partial), f'optimizer.type is not supported anymore, please use class path like "torch.optim.AdamW".'
-            weight_decay = cfg_opt_pt.keywords.get('weight_decay', None)
-            if weight_decay is not None:
-                for param in parameters_pt:
-                    param['weight_decay'] = weight_decay
-
-            parameters += parameters_pt
-            warnings.warn('deepspeed dose not support multi optimizer and lr_scheduler. optimizer_pt and scheduler_pt will not work.')
-
-        if len(parameters)>0:
-            cfg_opt = self.cfgs.train.optimizer
-            if self.cfgs.train.scale_lr:
-                self.scale_lr(parameters)
-            assert isinstance(cfg_opt, partial), f'optimizer.type is not supported anymore, please use class path like "torch.optim.AdamW".'
-            self.optimizer = cfg_opt(params=parameters)
-            self.lr_scheduler = get_scheduler(self.cfgs.train.scheduler, self.optimizer)
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Stable Diffusion Training')
-    parser.add_argument('--cfg', type=str, default='cfg/train/demo.yaml')
-    args, cfg_args = parser.parse_known_args()
-
-    conf = load_config_with_cli(args.cfg, args_list=cfg_args)  # skip --cfg
-    trainer = TrainerDeepSpeed(conf)
-    trainer.train()