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

hcpdiff/tools/lora_convert.py

@@ -3,14 +3,14 @@ import os.path
 from typing import List
 import math
 
-from hcpdiff.ckpt_manager import auto_manager
+from rainbowneko.ckpt_manager import auto_ckpt_loader, NekoModelSaver
 
 class LoraConverter:
     com_name_unet = ['down_blocks', 'up_blocks', 'mid_block', 'transformer_blocks', 'to_q', 'to_k', 'to_v', 'to_out', 'proj_in', 'proj_out', 'input_blocks', 'middle_block', 'output_blocks']
     com_name_TE = ['self_attn', 'q_proj', 'v_proj', 'k_proj', 'out_proj', 'text_model']
     prefix_unet = 'lora_unet_'
     prefix_TE = 'lora_te_'
-    prefix_TE_xl_clip_B = 'lora_te1_'
+    prefix_TE_xl_clip_L = 'lora_te1_'
     prefix_TE_xl_clip_bigG = 'lora_te2_'
 
     lora_w_map = {'lora_down.weight': 'W_down', 'lora_up.weight':'W_up'}
@@ -25,14 +25,14 @@ class LoraConverter:
             sd_TE = self.convert_from_webui_(state, prefix=self.prefix_TE, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
         else:
             sd_unet = self.convert_from_webui_xl_unet_(state, prefix=self.prefix_unet, com_name=self.com_name_unet, com_name_tmp=self.com_name_unet_tmp)
-            sd_TE = self.convert_from_webui_xl_te_(state, prefix=self.prefix_TE_xl_clip_B, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
+            sd_TE = self.convert_from_webui_xl_te_(state, prefix=self.prefix_TE_xl_clip_L, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
             sd_TE2 = self.convert_from_webui_xl_te_(state, prefix=self.prefix_TE_xl_clip_bigG, com_name=self.com_name_TE, com_name_tmp=self.com_name_TE_tmp)
             sd_TE.update(sd_TE2)
 
         if auto_scale_alpha:
             sd_unet = self.alpha_scale_from_webui(sd_unet)
             sd_TE = self.alpha_scale_from_webui(sd_TE)
-        return {'lora': sd_TE},  {'lora': sd_unet}
+        return {'plugin': sd_TE},  {'plugin': sd_unet}
 
     def convert_to_webui(self, sd_unet, sd_TE, auto_scale_alpha=False, sdxl=False):
         sd_unet = self.convert_to_webui_(sd_unet, prefix=self.prefix_unet)
@@ -72,7 +72,7 @@ class LoraConverter:
 
             sd_covert[f"{prefix}{model_k.replace('.', '_')}.{lora_k}"] = v
         return sd_covert
-    
+
     def convert_to_webui_xl_(self, state, prefix):
         sd_convert = {}
         for k, v in state.items():
@@ -87,7 +87,7 @@ class LoraConverter:
 
             new_k = f"{prefix}{model_k.replace('.', '_')}.{lora_k}"
             if 'clip' in new_k:
-                new_k = new_k.replace('_clip_B', '1') if 'clip_B' in new_k else new_k.replace('_clip_bigG', '2')
+                new_k = new_k.replace('_clip_L', '1') if 'clip_L' in new_k else new_k.replace('_clip_bigG', '2')
             sd_convert[new_k] = v
         return sd_convert
     
@@ -100,7 +100,7 @@ class LoraConverter:
             model_k, lora_k = k[prefix_len:].split('.', 1)
             model_k = self.replace_all(model_k, com_name, com_name_tmp).replace('_', '.').replace('%', '_')
             if prefix == 'lora_te1_':
-                model_k = f'clip_B.{model_k}'
+                model_k = f'clip_L.{model_k}'
             else:
                 model_k = f'clip_bigG.{model_k}'
 
@@ -221,23 +221,27 @@ if __name__ == '__main__':
 
     # load lora model
     print('convert lora model')
-    ckpt_manager = auto_manager(args.lora_path)
+    ckpt_loader = auto_ckpt_loader(args.lora_path)
+    ckpt_saver = NekoModelSaver(
+        format=ckpt_loader.format,
+        source=ckpt_loader.source,
+    )
 
     if args.from_webui:
-        state = ckpt_manager.load_ckpt(args.lora_path)
+        state = ckpt_loader.load(args.lora_path)
         # convert the weight name
         sd_TE, sd_unet = converter.convert_from_webui(state, auto_scale_alpha=args.auto_scale_alpha, sdxl=args.sdxl)
         # wegiht save
         os.makedirs(args.dump_path, exist_ok=True)
         TE_path = os.path.join(args.dump_path, 'TE-'+lora_name)
         unet_path = os.path.join(args.dump_path, 'unet-'+lora_name)
-        ckpt_manager._save_ckpt(sd_TE, save_path=TE_path)
-        ckpt_manager._save_ckpt(sd_unet, save_path=unet_path)
+        ckpt_saver.save(sd_TE, TE_path)
+        ckpt_saver.save(sd_unet, unet_path)
         print('save text encoder lora to:', TE_path)
         print('save unet lora to:', unet_path)
     elif args.to_webui:
-        sd_unet = ckpt_manager.load_ckpt(args.lora_path)
-        sd_TE = ckpt_manager.load_ckpt(args.lora_path_TE)
-        state = converter.convert_to_webui(sd_unet['lora'], sd_TE['lora'], auto_scale_alpha=args.auto_scale_alpha, sdxl=args.sdxl)
-        ckpt_manager._save_ckpt(state, save_path=args.dump_path)
+        sd_unet = ckpt_loader.load(args.lora_path)
+        sd_TE = ckpt_loader.load(args.lora_path_TE) if args.lora_path_TE else {'base':{}}
+        state = converter.convert_to_webui(sd_unet['base'], sd_TE['base'], auto_scale_alpha=args.auto_scale_alpha, sdxl=args.sdxl)
+        ckpt_saver.save(state, args.dump_path)
         print('save lora to:', args.dump_path)

hcpdiff/tools/save_model.py

@@ -0,0 +1,12 @@
+from diffusers import DiffusionPipeline
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("model", default=None, type=str)
+parser.add_argument("output", default=None, type=str)
+args = parser.parse_args()
+
+pipe = DiffusionPipeline.from_pretrained(args.model, safety_checker=None, requires_safety_checker=False,
+                                        resume_download=True)
+
+pipe.save_pretrained(args.output)

hcpdiff/tools/sd2diffusers.py

@@ -211,7 +211,7 @@ def sd_vae_to_diffuser(args):
 def convert_ckpt(args):
     pipe = load_sd_ckpt(
         args.checkpoint_path,
-        original_config_file=args.original_config_file,
+        config_files={'v1': args.original_config_file},
         image_size=args.image_size,
         prediction_type=args.prediction_type,
         model_type=args.pipeline_type,

hcpdiff/train_colo.py

@@ -23,7 +23,7 @@ from colossalai.utils.model.colo_init_context import ColoInitContext
 from colossalai.utils.model.colo_init_context import _convert_to_coloparam
 from colossalai.tensor import ColoParameter
 
-from hcpdiff.train_ac import Trainer, get_scheduler, ModelEMA
+from hcpdiff.train_ac_old import Trainer, get_scheduler, ModelEMA
 from diffusers import UNet2DConditionModel
 from hcpdiff.utils.colo_utils import gemini_zero_dpp, GeminiAdamOptimizerP
 from hcpdiff.utils.utils import load_config_with_cli

hcpdiff/train_deepspeed.py

@@ -7,7 +7,7 @@ from functools import partial
 import torch
 
 from hcpdiff.ckpt_manager import CkptManagerPKL, CkptManagerSafe
-from hcpdiff.train_ac import Trainer, load_config_with_cli
+from hcpdiff.train_ac_old import Trainer, load_config_with_cli
 from hcpdiff.utils.net_utils import get_scheduler
 
 class TrainerDeepSpeed(Trainer):

hcpdiff/trainer_ac.py

@@ -0,0 +1,79 @@
+import argparse
+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 hcpdiff.ckpt_manager import EmbFormat
+
+class HCPTrainer(Trainer):
+    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()
+
+    def get_param_group_train(self):
+        train_params = super().get_param_group_train()
+
+        # For prompt-tuning
+        if self.cfgs.emb_pt is None:
+            train_params_emb, self.train_pts = [], {}
+        else:
+            from hcpdiff.parser import CfgEmbPTParser
+            self.cfgs.emb_pt: CfgEmbPTParser
+
+            train_params_emb, self.train_pts = self.cfgs.emb_pt.get_params_group(self.model_wrapper)
+            self.emb_format = EmbFormat()
+        train_params += train_params_emb
+        return train_params
+
+    @property
+    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_ema=getattr(self, "ema_model", None),
+            name_template=f'{{}}-{self.real_step}',
+        )
+
+        self.loggers.info(f"Saved state, step: {self.real_step}")
+
+def hcp_train():
+    import subprocess
+    parser = argparse.ArgumentParser(description='HCP-Diffusion Launcher')
+    parser.add_argument('--launch_cfg', type=str, default='cfgs/launcher/multi.yaml')
+    args, train_args = parser.parse_known_args()
+
+    subprocess.run(["accelerate", "launch", '--config_file', args.launch_cfg, "-m",
+                    "hcpdiff.trainer_ac"] + train_args, check=True)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="HCP Diffusion Trainer")
+    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 = HCPTrainer(parser, conf)
+    trainer.train()

hcpdiff/trainer_ac_single.py

@@ -0,0 +1,31 @@
+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):
+    pass
+
+def hcp_train():
+    import subprocess
+    parser = argparse.ArgumentParser(description='HCP-Diffusion Launcher')
+    parser.add_argument('--launch_cfg', type=str, default='cfgs/launcher/single.yaml')
+    args, train_args = parser.parse_known_args()
+
+    subprocess.run(["accelerate", "launch", '--config_file', args.launch_cfg, "-m",
+                    "hcpdiff.trainer_ac_single"] + train_args, check=True)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='HCP Diffusion Trainer')
+    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 = HCPTrainerSingleCard(parser, conf)
+    trainer.train()

hcpdiff/utils/__init__.py

@@ -1,4 +1,2 @@
 from .utils import *
-from .img_size_tool import get_image_size
-from .cfg_resolvers import *
 from .net_utils import *