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

hcpdiff/ckpt_manager/__init__.py

@@ -1,4 +1,4 @@
 from .format import EmbFormat, DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSDXLFormat, \
-    OfficialSD15Format
+    OfficialSD15Format, LoraWebuiFormat
 from .ckpt import EmbSaver, easy_emb_saver
 from .loader import HCPLoraLoader

hcpdiff/ckpt_manager/format/lora_webui.py

@@ -2,7 +2,7 @@ import math
 import re
 from typing import List, Dict, Any
 
-from rainbowneko.ckpt_manager.format import CkptFormat
+from rainbowneko.ckpt_manager.format import CkptFormat, SafeTensorFormat
 from torch.serialization import FILE_LIKE
 
 class LoraConverter:
@@ -36,7 +36,12 @@ class LoraConverter:
         if auto_scale_alpha:
             sd_unet = self.alpha_scale_from_webui(sd_unet)
             sd_TE = self.alpha_scale_from_webui(sd_TE)
-        return {'plugin':sd_TE}, {'plugin':sd_unet}
+
+        sd = {
+            **{f'denoiser.{k}':v for k,v in sd_unet.items()},
+            **{f'TE.{k}':v for k,v in sd_TE.items()},
+        }
+        return {'base': sd}
 
     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)
@@ -207,9 +212,12 @@ class LoraConverter:
         return state
 
 class LoraWebuiFormat(CkptFormat):
-    def __init__(self, format, auto_scale_alpha=False):
+    def __init__(self, format=None, auto_scale_alpha=False):
         self.converter = LoraConverter()
         self.auto_scale_alpha = auto_scale_alpha
+
+        if format is None:
+            format = SafeTensorFormat()
         self.format = format
 
     def save_ckpt(self, sd_model: Dict[str, Any], save_f: FILE_LIKE):
@@ -240,5 +248,5 @@ class LoraWebuiFormat(CkptFormat):
                 sdxl = True
                 break
 
-        sd_TE, sd_unet = self.converter.convert_from_webui(sd_webui, auto_scale_alpha=self.auto_scale_alpha, sdxl=sdxl)
-        return sd_TE, sd_unet
+        sd_all = self.converter.convert_from_webui(sd_webui, auto_scale_alpha=self.auto_scale_alpha, sdxl=sdxl)
+        return sd_all

hcpdiff/data/__init__.py

@@ -1,4 +1,4 @@
 from .dataset import TextImagePairDataset
-from .source import Text2ImageSource, Text2ImageLossMapSource, Text2ImageCondSource, T2IFolderClassSource
-from .handler import StableDiffusionHandler, LossMapHandler, DiffusionImageHandler
+from .source import Text2ImageSource, Text2ImageLossMapSource, Text2ImageCondSource, T2IFolderClassSource, TextSource
+from .handler import StableDiffusionHandler, LossMapHandler, DiffusionImageHandler, DiffusionTextHandler
 from .cache import VaeCache

hcpdiff/data/handler/__init__.py

@@ -1,3 +1,3 @@
-from .diffusion import StableDiffusionHandler, DiffusionImageHandler, LossMapHandler
+from .diffusion import StableDiffusionHandler, DiffusionImageHandler, LossMapHandler, DiffusionTextHandler
 from .text import TokenizeHandler, TagEraseHandler, TagDropoutHandler, TagShuffleHandler, TemplateFillHandler
 from .controlnet import ControlNetHandler

hcpdiff/data/handler/diffusion.py

@@ -49,14 +49,11 @@ class DiffusionImageHandler(DataHandler):
         else:
             return self.handlers(dict(image=image, image_size=image_size))
 
-class StableDiffusionHandler(DataHandler):
-    def __init__(self, bucket, encoder_attention_mask=False, key_map_in=('image -> image', 'image_size -> image_size', 'prompt -> prompt'),
-                 key_map_out=('image -> image', 'coord -> coord', 'prompt -> prompt'),
-                 erase=0.15, dropout=0.0, shuffle=0.0, word_names={}, tokenize=True):
+class DiffusionTextHandler(DataHandler):
+    def __init__(self, encoder_attention_mask=False, erase=0.0, dropout=0.0, shuffle=0.0, word_names={}, tokenize=True,
+                 key_map_in=('prompt -> prompt', ), key_map_out=('prompt -> prompt', )):
         super().__init__(key_map_in, key_map_out)
 
-        self.image_handlers = DiffusionImageHandler(bucket)
-
         text_handlers = {}
         if dropout>0:
             text_handlers['dropout'] = TagDropoutHandler(p=dropout)
@@ -67,7 +64,20 @@ class StableDiffusionHandler(DataHandler):
         text_handlers['fill'] = TemplateFillHandler(word_names)
         if tokenize:
             text_handlers['tokenize'] = TokenizeHandler(encoder_attention_mask)
-        self.text_handlers = HandlerChain(**text_handlers)
+        self.handlers = HandlerChain(**text_handlers)
+
+    def handle(self, prompt: Union[str, Dict[str, str]]):
+        return self.handlers(dict(prompt=prompt))
+
+class StableDiffusionHandler(DataHandler):
+    def __init__(self, bucket, encoder_attention_mask=False, key_map_in=('image -> image', 'image_size -> image_size', 'prompt -> prompt'),
+                 key_map_out=('image -> image', 'coord -> coord', 'prompt -> prompt'),
+                 erase=0.0, dropout=0.0, shuffle=0.0, word_names={}, tokenize=True):
+        super().__init__(key_map_in, key_map_out)
+
+        self.image_handlers = DiffusionImageHandler(bucket)
+        self.text_handlers = DiffusionTextHandler(encoder_attention_mask=encoder_attention_mask, erase=erase, dropout=dropout, shuffle=shuffle,
+                                                  word_names=word_names, tokenize=tokenize)
 
     def handle(self, image: Image.Image, image_size: np.ndarray[int], prompt: str):
         return dict(**self.image_handlers(dict(image=image, image_size=image_size)), **self.text_handlers(dict(prompt=prompt)))

hcpdiff/data/source/__init__.py

@@ -1,3 +1,4 @@
 from .text2img import Text2ImageSource, Text2ImageLossMapSource
 from .text2img_cond import Text2ImageCondSource
-from .folder_class import T2IFolderClassSource
+from .folder_class import T2IFolderClassSource
+from .text import TextSource

hcpdiff/data/source/text.py

@@ -0,0 +1,40 @@
+from rainbowneko.data import UnLabelSource, DataSource
+from rainbowneko.data.label_loader import BaseLabelLoader, auto_label_loader
+from typing import Union, Dict, Any
+import random
+
+class TextSource(DataSource):
+    def __init__(self, label_file, prompt_template=None, repeat=1, **kwargs):
+        super().__init__(repeat=repeat)
+        self.label_file = label_file
+        self.label_dict = self._load_label_data(label_file)
+        self.img_ids = self._load_img_ids(self.label_dict)
+        self.prompt_template = self.load_template(prompt_template)
+
+    def _load_img_ids(self, label_dict):
+        return list(label_dict.keys()) * self.repeat
+
+    def _load_label_data(self, label_file: Union[str, BaseLabelLoader]):
+        if label_file is None:
+            return {}
+        elif isinstance(label_file, str):
+            return auto_label_loader(label_file).load()
+        else:
+            return label_file.load()
+
+    def load_template(self, template_file):
+        if template_file is None:
+            return ['{caption}']
+        else:
+            with open(template_file, 'r', encoding='utf-8') as f:
+                return f.read().strip().split('\n')
+
+    def __getitem__(self, index) -> Dict[str, Any]:
+        img_name = self.img_ids[index]
+        return {
+            'id':img_name,
+            'prompt':{
+                'template':random.choice(self.prompt_template),
+                'caption':self.label_dict[img_name],
+            }
+        }

hcpdiff/data/source/text2img.py

@@ -25,7 +25,7 @@ class Text2ImageSource(ImageLabelSource):
 
     def __getitem__(self, index) -> Dict[str, Any]:
         img_name = self.img_ids[index]
-        path = os.path.join(self.img_root, img_name)
+        path = self.img_root / img_name
 
         return {
             'id':img_name,

hcpdiff/easy/cfg/__init__.py

@@ -1,3 +1,3 @@
 from .sd15_train import SD15_lora_train, cfg_data_SD_ARB, cfg_data_SD_resize_crop, SD15_finetuning
 from .sdxl_train import SDXL_lora_train, SDXL_finetuning
-from .t2i import SD15_t2i, SDXL_t2i, SDXL_t2i_lora, SD15_t2i_lora
+from .t2i import SD15_t2i, SDXL_t2i, SDXL_t2i_lora, SD15_t2i_lora, SDXL_t2i_parts, SD15_t2i_parts

hcpdiff/easy/cfg/sd15_train.py

@@ -1,9 +1,10 @@
 import torch
-from rainbowneko.ckpt_manager import ckpt_saver, LAYERS_TRAINABLE, plugin_saver
+from rainbowneko.ckpt_manager import ckpt_saver, LAYERS_TRAINABLE, NekoPluginSaver, SafeTensorFormat
 from rainbowneko.data import RatioBucket, FixedBucket
 from rainbowneko.parser import CfgWDPluginParser, neko_cfg, CfgWDModelParser, disable_neko_cfg
 from rainbowneko.utils import ConstantLR, Path_Like
 
+from hcpdiff.ckpt_manager import LoraWebuiFormat
 from hcpdiff.data import TextImagePairDataset, Text2ImageSource, StableDiffusionHandler
 from hcpdiff.data import VaeCache
 from hcpdiff.easy import SD15_auto_loader
@@ -46,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,
             ),
@@ -69,7 +70,7 @@ def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
 @neko_cfg
 def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int = 200, lr: float = 1e-4, rank: int = 4, alpha: float = None,
                     clip_skip: int = 0, with_conv: bool = False, dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0,
-                    name: str = 'SD15'):
+                    name: str = 'SD15', save_webui_format=False):
     with disable_neko_cfg:
         if alpha is None:
             alpha = rank
@@ -95,6 +96,11 @@ def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
     else:
         optimizer = torch.optim.AdamW(_partial_=True, betas=(0.9, 0.99))
 
+    if save_webui_format:
+        lora_format = LoraWebuiFormat()
+    else:
+        lora_format = SafeTensorFormat()
+
     from cfgs.train.py.examples import SD_FT
 
     return dict(
@@ -114,8 +120,8 @@ def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
 
         ckpt_saver=dict(
             _replace_ = True,
-            lora_unet=plugin_saver(
-                ckpt_type='safetensors',
+            lora_unet=NekoPluginSaver(
+                format=lora_format,
                 target_plugin='lora1',
             )
         ),
@@ -126,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

@@ -1,11 +1,12 @@
 import torch
-from rainbowneko.ckpt_manager import ckpt_saver, plugin_saver, LAYERS_TRAINABLE
+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.easy import SDXL_auto_loader
 from hcpdiff.models import SDXLWrapper
 from hcpdiff.models.lora_layers_patch import LoraLayer
+from hcpdiff.ckpt_manager import LoraWebuiFormat
 
 @neko_cfg
 def SDXL_finetuning(base_model: str, train_steps: int, dataset, save_step: int = 500, lr: float = 1e-5,
@@ -43,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,
             ),
@@ -64,7 +65,8 @@ 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 = 'SD15'):
+                    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
@@ -90,6 +92,11 @@ def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
     else:
         optimizer = torch.optim.AdamW(_partial_=True, betas=(0.9, 0.99))
 
+    if save_webui_format:
+        lora_format = LoraWebuiFormat()
+    else:
+        lora_format = SafeTensorFormat()
+
     from cfgs.train.py.examples import SD_FT
 
     return dict(
@@ -109,8 +116,8 @@ def SDXL_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
 
         ckpt_saver=dict(
             _replace_ = True,
-            lora_unet=plugin_saver(
-                ckpt_type='safetensors',
+            lora_unet=NekoPluginSaver(
+                format=lora_format,
                 target_plugin='lora1',
             )
         ),
@@ -121,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/easy/cfg/t2i.py

@@ -1,5 +1,6 @@
 import torch
 from rainbowneko.infer.workflow import (Actions, PrepareAction, LoopAction, LoadModelAction)
+from rainbowneko.ckpt_manager import NekoModelLoader
 from rainbowneko.parser import neko_cfg, disable_neko_cfg
 from typing import Union, List
 
@@ -25,6 +26,29 @@ def build_model(pretrained_model='ckpts/any5', noise_sampler=Diffusers_SD.dpmpp_
         ),
     ])
 
+@neko_cfg
+def load_parts(info: List[str]) -> Actions:
+    acts = []
+    for i, path in enumerate(info):
+        part_unet = LoadModelAction(cfg={
+            f'part_unet_{i}':NekoModelLoader(
+                path=path,
+                state_prefix='denoiser.'
+            )
+        }, key_map_in=('denoiser -> model', 'in_preview -> in_preview'))
+        part_TE = LoadModelAction(cfg={
+            f'part_TE_{i}':NekoModelLoader(
+                path=path,
+                state_prefix='TE.',
+            )
+        }, key_map_in=('TE -> model', 'in_preview -> in_preview'))
+
+        with disable_neko_cfg:
+            acts.append(part_unet)
+            acts.append(part_TE)
+
+    return Actions(acts)
+
 @neko_cfg
 def load_lora(info: List[List]) -> Actions:
     lora_acts = []
@@ -37,7 +61,7 @@ def load_lora(info: List[List]) -> Actions:
             )
         }, key_map_in=('denoiser -> model', 'in_preview -> in_preview'))
         lora_TE = LoadModelAction(cfg={
-            f'lora_unet_{i}':HCPLoraLoader(
+            f'lora_TE_{i}':HCPLoraLoader(
                 path=item[0],
                 state_prefix='TE.',
                 alpha=item[1],
@@ -59,9 +83,9 @@ def optimize_model() -> Actions:
     ])
 
 @neko_cfg
-def text(prompt, negative_prompt=negative_prompt, bs=4) -> Actions:
+def text(prompt, negative_prompt=negative_prompt, bs=4, N_repeats=1, layer_skip=1) -> Actions:
     return Actions([
-        TextHookAction(N_repeats=1, layer_skip=1),
+        TextHookAction(N_repeats=N_repeats, layer_skip=layer_skip),
         AttnMultTextEncodeAction(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -84,9 +108,9 @@ def build_model_SDXL(pretrained_model='ckpts/any5', noise_sampler=Diffusers_SD.d
     ])
 
 @neko_cfg
-def text_SDXL(prompt, negative_prompt=negative_prompt, bs=4) -> Actions:
+def text_SDXL(prompt, negative_prompt=negative_prompt, bs=4, N_repeats=1, layer_skip=1) -> Actions:
     return Actions([
-        TextHookAction(N_repeats=1, layer_skip=1, TE_final_norm=False),
+        TextHookAction(N_repeats=N_repeats, layer_skip=layer_skip, TE_final_norm=False),
         AttnMultTextEncodeAction(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -128,11 +152,24 @@ def resize(width=1024, height=1024):
 
 @neko_cfg
 def SD15_t2i(pretrained_model, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4, width=512, height=512,
-             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/', N_repeats=1, layer_skip=1):
+    return dict(workflow=Actions(actions=[
+        build_model(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        optimize_model(),
+        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs, N_repeats=N_repeats, layer_skip=layer_skip),
+        config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
+        diffusion(guidance_scale=guidance_scale),
+        decode(save_root=save_root)
+    ]))
+
+@neko_cfg
+def SD15_t2i_parts(pretrained_model, parts, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4, width=512, height=512,
+             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/', N_repeats=1, layer_skip=1):
     return dict(workflow=Actions(actions=[
         build_model(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        load_parts(parts),
         optimize_model(),
-        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs, N_repeats=N_repeats, layer_skip=layer_skip),
         config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
         diffusion(guidance_scale=guidance_scale),
         decode(save_root=save_root)
@@ -140,12 +177,12 @@ def SD15_t2i(pretrained_model, prompt, negative_prompt=negative_prompt, noise_sa
 
 @neko_cfg
 def SD15_t2i_lora(pretrained_model, lora_info, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4,
-                  width=512, height=512, seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+                  width=512, height=512, seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/', N_repeats=1, layer_skip=1):
     return dict(workflow=Actions(actions=[
         build_model(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
         load_lora(info=lora_info),
         optimize_model(),
-        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        text(prompt=prompt, negative_prompt=negative_prompt, bs=bs, N_repeats=N_repeats, layer_skip=layer_skip),
         config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
         diffusion(guidance_scale=guidance_scale),
         decode(save_root=save_root)
@@ -153,24 +190,38 @@ def SD15_t2i_lora(pretrained_model, lora_info, prompt, negative_prompt=negative_
 
 @neko_cfg
 def SDXL_t2i(pretrained_model, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4, width=1024, height=1024,
-             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/', N_repeats=1, layer_skip=1):
     return dict(workflow=Actions(actions=[
         build_model_SDXL(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
         optimize_model(),
-        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs, N_repeats=N_repeats, layer_skip=layer_skip),
         config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
         diffusion(guidance_scale=guidance_scale),
         decode(save_root=save_root)
     ]))
 
+@neko_cfg
+def SDXL_t2i_parts(pretrained_model, parts, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4, width=1024, height=1024,
+             seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/', N_repeats=1, layer_skip=1):
+    return dict(workflow=Actions(actions=[
+        build_model_SDXL(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
+        load_parts(parts),
+        optimize_model(),
+        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs, N_repeats=N_repeats, layer_skip=layer_skip),
+        config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
+        diffusion(guidance_scale=guidance_scale),
+        decode(save_root=save_root)
+    ]))
+
+
 @neko_cfg
 def SDXL_t2i_lora(pretrained_model, lora_info, prompt, negative_prompt=negative_prompt, noise_sampler=Diffusers_SD.dpmpp_2m_karras, bs=4,
-                  width=1024, height=1024, seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/'):
+                  width=1024, height=1024, seed=None, N_steps=20, guidance_scale=7.0, save_root='output_pipe/', N_repeats=1, layer_skip=1):
     return dict(workflow=Actions(actions=[
         build_model_SDXL(pretrained_model=pretrained_model, noise_sampler=noise_sampler),
         load_lora(info=lora_info),
         optimize_model(),
-        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs),
+        text_SDXL(prompt=prompt, negative_prompt=negative_prompt, bs=bs, N_repeats=N_repeats, layer_skip=layer_skip),
         config_diffusion(width=width, height=height, seed=seed, N_steps=N_steps),
         diffusion(guidance_scale=guidance_scale),
         decode(save_root=save_root)

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/workflow/diffusion.py

@@ -32,14 +32,15 @@ class SeedAction(BasicAction):
         self.seed = seed
         self.bs = bs
 
-    def forward(self, device, gen_step=0, **states):
+    def forward(self, device, seed=None, **states):
         bs = states['prompt_embeds'].shape[0]//2 if 'prompt_embeds' in states else self.bs
-        if self.seed is None:
+        seed = seed or self.seed
+        if seed is None:
             seeds = [None]*bs
-        elif isinstance(self.seed, int):
-            seeds = list(range(self.seed+gen_step*bs, self.seed+(gen_step+1)*bs))
+        elif isinstance(seed, int):
+            seeds = list(range(seed, seed+bs))
         else:
-            seeds = self.seed
+            seeds = seed
         seeds = [s or random.randint(0, 1 << 30) for s in seeds]
 
         G = prepare_seed(seeds, device=device)

hcpdiff/workflow/text.py

@@ -48,18 +48,9 @@ class TextEncodeAction(BasicAction):
         self.negative_prompt = negative_prompt
         self.bs = bs
 
-    def forward(self, te_hook, TE, dtype: str, device, amp=None, gen_step=None, prompt_all=None, negative_prompt_all=None, model_offload=False,
-                **states):
-        prompt_all = prompt_all or self.prompt
-        negative_prompt_all = negative_prompt_all or self.negative_prompt
-
-        if gen_step is not None:
-            idx = (gen_step*self.bs)%len(prompt_all)
-            prompt = prompt_all[idx:idx+self.bs]
-            negative_prompt = negative_prompt_all[idx:idx+self.bs]
-        else:
-            prompt = prompt_all
-            negative_prompt = negative_prompt_all
+    def forward(self, te_hook, TE, dtype: str, device, amp=None, prompt=None, negative_prompt=None, model_offload=False, **states):
+        prompt = prompt or self.prompt
+        negative_prompt = negative_prompt or self.negative_prompt
 
         if model_offload:
             to_cuda(TE)
@@ -78,19 +69,9 @@ class TextEncodeAction(BasicAction):
             'pooled_output':pooled_output}
 
 class AttnMultTextEncodeAction(TextEncodeAction):
-
-    def forward(self, te_hook, token_ex, TE, dtype: str, device, amp=None, gen_step=None, prompt_all=None, negative_prompt_all=None,
-                model_offload=False, **states):
-        prompt_all = prompt_all if prompt_all is not None else self.prompt
-        negative_prompt_all = negative_prompt_all if negative_prompt_all is not None else self.negative_prompt
-
-        if gen_step is not None:
-            idx = (gen_step*self.bs)%len(prompt_all)
-            prompt = prompt_all[idx:idx+self.bs]
-            negative_prompt = negative_prompt_all[idx:idx+self.bs]
-        else:
-            prompt = prompt_all
-            negative_prompt = negative_prompt_all
+    def forward(self, te_hook, token_ex, TE, dtype: str, device, amp=None, prompt=None, negative_prompt=None, model_offload=False, **states):
+        prompt = prompt or self.prompt
+        negative_prompt = negative_prompt or self.negative_prompt
 
         if model_offload:
             to_cuda(TE)

{hcpdiff-2.1.dist-info → hcpdiff-2.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hcpdiff
-Version: 2.1
+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
@@ -65,6 +65,8 @@ Compared to the original DreamArtist, it offers better stability, image quality,
 
 ## Installation
 
+Install [pytorch](https://pytorch.org/)
+
 Install via pip:
 
 ```bash
@@ -205,6 +207,18 @@ After parsing, the framework will instantiate the components accordingly. This m
 | CCIP Score       | 🚧 In Development                        |
 | Corrupt Score    | 🚧 In Development                        |
 
+---
+
+### ⚡️ Image Generation
+
+| 功能                           | 描述/支持情况                            |
+|------------------------------|------------------------------------|
+| Batch Generation             | ✅ Supported                   |
+| Generate from Prompt Dataset | ✅ Supported                               |
+| Image to Image               | ✅ Supported                               |
+| Inpaint                      | ✅ Supported                               |
+| Token Weight                 | ✅ Supported |
+
 </details>
 
 ---
@@ -248,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.1.dist-info → hcpdiff-2.2.1.dist-info}/RECORD

@@ -1,27 +1,28 @@
 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/ckpt_manager/__init__.py,sha256=LfMwz9R4jV4xpiSFt5vhpwaF7-8UHEZ_iDoW-3QGvt0,239
+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/format/__init__.py,sha256=a3cdKkOTDgdVbDQwSC4mlxOigjX2hBvRb5_X7E3TQWs,237
 hcpdiff/ckpt_manager/format/diffusers.py,sha256=T81WN95Nj1il9DfQp9iioVn0uqFEWOlmdIYs2beNOFU,3769
 hcpdiff/ckpt_manager/format/emb.py,sha256=FrqfTfJ8H7f0Zw17NTWCP2AJtpsJI5oXR5IAd4NekhU,680
-hcpdiff/ckpt_manager/format/lora_webui.py,sha256=j7SpXnSx_Ys8tnWBgojuB1HEJIm46lhCBuNNYLhaF9w,9824
+hcpdiff/ckpt_manager/format/lora_webui.py,sha256=4y_T9RdmFTxWzsXd8guNjCiukmyILa5j4MPrhVIL4Qk,10017
 hcpdiff/ckpt_manager/format/sd_single.py,sha256=LpCAL_7nAVooCHTFznVVsNMku1G3C77NBORxxr8GDtQ,2328
-hcpdiff/data/__init__.py,sha256=-z47HsEQSubc-AfriVComMACbQXlXTWAKMOPBkATHxA,258
+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
 hcpdiff/data/cache/vae.py,sha256=gB89zs4CdNlvukDXhVYU9QZrY6VTFUWfzjeF2psNQ50,4070
-hcpdiff/data/handler/__init__.py,sha256=D1HyqY0qfrUHgf25itpYj57JUvgn06G6EQ9d2vRRtys,236
+hcpdiff/data/handler/__init__.py,sha256=G8ZTQF91ilkTRmUoWdmAissTSZ7fvNUpm_hBYmXKTtk,258
 hcpdiff/data/handler/controlnet.py,sha256=bRDMD9BP8-VaG5VrxzvcFKfkqeTbChNfrJSZ3vXbQgY,658
-hcpdiff/data/handler/diffusion.py,sha256=8n60UYdGNR08xw45HoI4EB5AaIui03tSGNDfjazO-5w,3516
+hcpdiff/data/handler/diffusion.py,sha256=S-_7o5Z1tm6LmRZVZs21rbJC7iUoq0tHOsSjKK6geVk,4156
 hcpdiff/data/handler/text.py,sha256=gOzqB2oEkEUbiuy0kZWduo0c-w4Buu60KI6q6Nyl3aM,4208
-hcpdiff/data/source/__init__.py,sha256=AB1VicA272KjTm-Q5L6XvDM8CLQhVPylAPuPMtpfw4g,158
+hcpdiff/data/source/__init__.py,sha256=265M8qfWNUE4SKX0pdXhLYjCnCuae5YE4bfZpO-ydXc,187
 hcpdiff/data/source/folder_class.py,sha256=bs4qPMTzwcnT6ZFlT3tpi9sclsRF9a2MBA1pQD-9EYs,961
-hcpdiff/data/source/text2img.py,sha256=MWXqAEbzmK6pkBY40t9u37ngY25mgdKQ2idwNld8-bo,1826
+hcpdiff/data/source/text.py,sha256=VgI5Ouq986Yy1jwD2fZ9iBlsRciPCeARZmOPEZIcaQY,1468
+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
@@ -38,10 +39,10 @@ hcpdiff/diffusion/sampler/sigma_scheduler/ddpm.py,sha256=2PMIpg2K6CVoxew1y1pIqvC
 hcpdiff/diffusion/sampler/sigma_scheduler/edm.py,sha256=fOPB3lgnS9uVo4oW26Fur_nc8X_wQ6mmUcbkKhnoQjs,1900
 hcpdiff/easy/__init__.py,sha256=-emoyCOZlLCu3KNMI8L4qapUEtEYFSoiGU6-rKv1at4,149
 hcpdiff/easy/sampler.py,sha256=dQSBkeGh71O0DAmZLhTHTbk1bY7XzyUCeW1oJO14A4I,1250
-hcpdiff/easy/cfg/__init__.py,sha256=aVDEDPxHdX5n-aFkP_4ic8ZhQfSeKu8lZOkgW_4m398,221
-hcpdiff/easy/cfg/sd15_train.py,sha256=LRCJLHNU0JEd1m3MC_NFWUCw5LmwztiLiJlV7u_DeKM,6493
-hcpdiff/easy/cfg/sdxl_train.py,sha256=R0wolSVOrRlI9A-vAfz592SzSnwuDd4ku1oc5yRKrfU,4038
-hcpdiff/easy/cfg/t2i.py,sha256=6Pyy4werXNalwoBBHVMBLBg67kMS85Heb7R3t26GJqQ,6871
+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/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
@@ -61,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
@@ -95,20 +96,20 @@ hcpdiff/utils/net_utils.py,sha256=gdwLYDNKV2t3SP0jBIO3d0HtY6E7jRaf_rmPT8gKZZE,97
 hcpdiff/utils/pipe_hook.py,sha256=-UDX3FtZGl-bxSk13gdbPXc1OvtbCcpk_fvKxLQo3Ag,31987
 hcpdiff/utils/utils.py,sha256=hZnZP1IETgVpScxES0yIuRfc34TnzvAqmgOTK_56ssw,4976
 hcpdiff/workflow/__init__.py,sha256=t7Zyc0XFORdNvcwHp9AsCtEkhJ3l7Hm41ugngIL0Sag,867
-hcpdiff/workflow/diffusion.py,sha256=yrl2cXE2d2FNeVzYZDRQNLjy5-QnVgOWioIHSmszk2Y,8662
+hcpdiff/workflow/diffusion.py,sha256=yzhqKA3019OPu1RKggrLoytMgm919qf6j9S85PYOwjQ,8644
 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/model.py,sha256=1gj5yOTefYTnGXVR6JPAfxIwuB69YwN6E-BontRcuyQ,2913
-hcpdiff/workflow/text.py,sha256=FSFUm_zEeZjMeg0qRXZAPplnJkg2pR_2FA3XljpoN2w,5110
+hcpdiff/workflow/text.py,sha256=Z__SJHZyuaKyzkYJ6rbiAzOGRiYcCjwCGeqfpP1Jo7o,4336
 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.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-hcpdiff-2.1.dist-info/METADATA,sha256=NpBZuj23d1gTKPQhJ0TBRV8QsfICa4LCGSk6PJNniSw,9248
-hcpdiff-2.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-hcpdiff-2.1.dist-info/entry_points.txt,sha256=86wPOMzsfWWflTJ-sQPLc7WG5Vtu0kGYBH9C_vR3ur8,207
-hcpdiff-2.1.dist-info/top_level.txt,sha256=shyf78x-HVgykYpsmY22mKG0xIc7Qk30fDMdavdYWQ8,8
-hcpdiff-2.1.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.1.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.1.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()