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

hcpdiff/__init__.py

@@ -1,4 +1,4 @@
-from .train_ac import Trainer
-from .train_ac_single import TrainerSingleCard
-from .visualizer import Visualizer
-from .visualizer_reloadable import VisualizerReloadable
+#from .train_ac_old import Trainer
+#from .train_ac_single import TrainerSingleCard
+# from .visualizer import Visualizer
+# from .visualizer_reloadable import VisualizerReloadable

hcpdiff/ckpt_manager/__init__.py

@@ -1,5 +1,4 @@
-from .ckpt_pkl import CkptManagerPKL
-from .ckpt_safetensor import CkptManagerSafe
-
-def auto_manager(ckpt_path:str):
-    return CkptManagerSafe() if ckpt_path.endswith('.safetensors') else CkptManagerPKL()
+from .format import EmbFormat, DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSDXLFormat, \
+    OfficialSD15Format, LoraWebuiFormat
+from .ckpt import EmbSaver, easy_emb_saver
+from .loader import HCPLoraLoader

hcpdiff/ckpt_manager/ckpt.py

@@ -0,0 +1,24 @@
+from rainbowneko.ckpt_manager import NekoSaver, CkptFormat, LocalCkptSource, PKLFormat
+from torch import nn
+from typing import Dict, Any
+
+class EmbSaver(NekoSaver):
+    def __init__(self, format: CkptFormat, source: LocalCkptSource, target_key='embs', prefix=None):
+        super().__init__(format, source)
+        self.target_key = target_key
+        self.prefix = prefix
+
+    def save_to(self, name, model: nn.Module, plugin_groups: Dict[str, Any], model_ema=None, exclude_key=None,
+                name_template=None):
+        train_pts = plugin_groups[self.target_key]
+        for pt_name, pt in train_pts.items():
+            self.save(pt_name, (pt_name, pt), prefix=self.prefix)
+            if name_template is not None:
+                pt_name = name_template.format(pt_name)
+                self.save(pt_name, (pt_name, pt), prefix=self.prefix)
+
+def easy_emb_saver():
+    return EmbSaver(
+        format=PKLFormat(),
+        source=LocalCkptSource(),
+    )

hcpdiff/ckpt_manager/format/__init__.py

@@ -0,0 +1,4 @@
+from .emb import EmbFormat
+from .diffusers import DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat
+from .sd_single import OfficialSD15Format, OfficialSDXLFormat
+from .lora_webui import LoraWebuiFormat

hcpdiff/ckpt_manager/format/diffusers.py

@@ -0,0 +1,59 @@
+import torch
+from diffusers import ModelMixin, AutoencoderKL, UNet2DConditionModel, PixArtTransformer2DModel
+from rainbowneko.ckpt_manager.format import CkptFormat
+from transformers import CLIPTextModel, AutoTokenizer, T5EncoderModel
+
+from hcpdiff.diffusion.sampler import DDPMSampler, DDPMDiscreteSigmaScheduler
+from hcpdiff.models.compose import SDXLTokenizer, SDXLTextEncoder
+
+class DiffusersModelFormat(CkptFormat):
+    def __init__(self, builder: ModelMixin):
+        self.builder = builder
+
+    def save_ckpt(self, sd_model: ModelMixin, save_f: str, **kwargs):
+        sd_model.save_pretrained(save_f)
+
+    def load_ckpt(self, ckpt_f: str, map_location="cpu", **kwargs):
+        self.builder.from_pretrained(ckpt_f, **kwargs)
+
+class DiffusersSD15Format(CkptFormat):
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        denoiser = denoiser or UNet2DConditionModel.from_pretrained(
+            pretrained_model, subfolder="unet", revision=revision, torch_dtype=dtype
+        )
+        vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
+        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+
+        TE = TE or CLIPTextModel.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
+        tokenizer = tokenizer or AutoTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
+
+        return dict(denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer)
+
+class DiffusersSDXLFormat(CkptFormat):
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        denoiser = denoiser or UNet2DConditionModel.from_pretrained(
+            pretrained_model, subfolder="unet", revision=revision, torch_dtype=dtype
+        )
+        vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
+        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+
+        TE = TE or SDXLTextEncoder.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
+        tokenizer = tokenizer or SDXLTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
+
+        return dict(denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer)
+
+class DiffusersPixArtFormat(CkptFormat):
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        denoiser = denoiser or PixArtTransformer2DModel.from_pretrained(
+            pretrained_model, subfolder="transformer", revision=revision, torch_dtype=dtype
+        )
+        vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
+        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+
+        TE = TE or T5EncoderModel.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
+        tokenizer = tokenizer or AutoTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
+
+        return dict(denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer)

hcpdiff/ckpt_manager/format/emb.py

@@ -0,0 +1,21 @@
+from typing import Tuple
+
+import torch
+from rainbowneko.ckpt_manager.format import CkptFormat
+from torch.serialization import FILE_LIKE
+
+class EmbFormat(CkptFormat):
+    EXT = 'pt'
+
+    def save_ckpt(self, sd_model: Tuple[str, torch.Tensor], save_f: FILE_LIKE):
+        name, emb = sd_model
+        torch.save({'string_to_param':{'*':emb}, 'name':name}, save_f)
+
+    def load_ckpt(self, ckpt_f: FILE_LIKE, map_location="cpu"):
+        state = torch.load(ckpt_f, map_location=map_location)
+        if 'string_to_param' in state:
+            emb = state['string_to_param']['*']
+        else:
+            emb = state['emb_params']
+        emb.requires_grad_(False)
+        return emb

hcpdiff/ckpt_manager/format/lora_webui.py

@@ -0,0 +1,252 @@
+import math
+import re
+from typing import List, Dict, Any
+
+from rainbowneko.ckpt_manager.format import CkptFormat, SafeTensorFormat
+from torch.serialization import FILE_LIKE
+
+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_L = 'lora_te1_'
+    prefix_TE_xl_clip_bigG = 'lora_te2_'
+
+    lora_w_map = {'lora_down.weight':'W_down', 'lora_up.weight':'W_up'}
+
+    def __init__(self):
+        self.com_name_unet_tmp = [x.replace('_', '%') for x in self.com_name_unet]
+        self.com_name_TE_tmp = [x.replace('_', '%') for x in self.com_name_TE]
+
+    def convert_from_webui(self, state, auto_scale_alpha=False, sdxl=False):
+        if not sdxl:
+            sd_unet = self.convert_from_webui_(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_(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_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)
+
+        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)
+        if sdxl:
+            sd_TE = self.convert_to_webui_xl_(sd_TE, prefix=self.prefix_TE)
+        else:
+            sd_TE = self.convert_to_webui_(sd_TE, prefix=self.prefix_TE)
+        sd_unet.update(sd_TE)
+        if auto_scale_alpha:
+            sd_unet = self.alpha_scale_to_webui(sd_unet)
+        return sd_unet
+
+    def convert_from_webui_(self, state, prefix, com_name, com_name_tmp):
+        state = {k:v for k, v in state.items() if k.startswith(prefix)}
+        prefix_len = len(prefix)
+        sd_covert = {}
+        for k, v in state.items():
+            model_k, lora_k = k[prefix_len:].split('.', 1)
+            model_k = self.replace_all(model_k, com_name, com_name_tmp).replace('_', '.').replace('%', '_')
+            if lora_k == 'alpha':
+                sd_covert[f'{model_k}.___.{lora_k}'] = v
+            else:
+                sd_covert[f'{model_k}.___.layer.{self.lora_w_map[lora_k]}'] = v
+        return sd_covert
+
+    def convert_to_webui_(self, state, prefix):
+        sd_covert = {}
+        for k, v in state.items():
+            if k.endswith('W_down'):
+                model_k, _ = k.split('.___.', 1)
+                lora_k = 'lora_down.weight'
+            elif k.endswith('W_up'):
+                model_k, _ = k.split('.___.', 1)
+                lora_k = 'lora_up.weight'
+            else:
+                model_k, lora_k = k.split('.___.', 1)
+
+            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():
+            if k.endswith('W_down'):
+                model_k, _ = k.split('.___.', 1)
+                lora_k = 'lora_down.weight'
+            elif k.endswith('W_up'):
+                model_k, _ = k.split('.___.', 1)
+                lora_k = 'lora_up.weight'
+            else:
+                model_k, lora_k = k.split('.___.', 1)
+
+            new_k = f"{prefix}{model_k.replace('.', '_')}.{lora_k}"
+            if 'clip' in new_k:
+                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
+
+    def convert_from_webui_xl_te_(self, state, prefix, com_name, com_name_tmp):
+        state = {k:v for k, v in state.items() if k.startswith(prefix)}
+        sd_covert = {}
+        prefix_len = len(prefix)
+
+        for k, v in state.items():
+            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_L.{model_k}'
+            else:
+                model_k = f'clip_bigG.{model_k}'
+
+            if lora_k == 'alpha':
+                sd_covert[f'{model_k}.___.{lora_k}'] = v
+            else:
+                sd_covert[f'{model_k}.___.layer.{self.lora_w_map[lora_k]}'] = v
+        return sd_covert
+
+    def convert_from_webui_xl_unet_(self, state, prefix, com_name, com_name_tmp):
+        # Down:
+        # 4 -> 1, 0  4 = 1 + 3 * 1 + 0
+        # 5 -> 1, 1  5 = 1 + 3 * 1 + 1
+        # 7 -> 2, 0  7 = 1 + 3 * 2 + 0
+        # 8 -> 2, 1  8 = 1 + 3 * 2 + 1
+
+        # Up
+        # 0 -> 0, 0  0 = 0 * 3 + 0
+        # 1 -> 0, 1  1 = 0 * 3 + 1
+        # 2 -> 0, 2  2 = 0 * 3 + 2
+        # 3 -> 1, 0  3 = 1 * 3 + 0
+        # 4 -> 1, 1  4 = 1 * 3 + 1
+        # 5 -> 1, 2  5 = 1 * 3 + 2
+
+        down = {
+            '4':[1, 0],
+            '5':[1, 1],
+            '7':[2, 0],
+            '8':[2, 1],
+        }
+        up = {
+            '0':[0, 0],
+            '1':[0, 1],
+            '2':[0, 2],
+            '3':[1, 0],
+            '4':[1, 1],
+            '5':[1, 2],
+        }
+
+        m = []
+
+        def match(key, regex_text):
+            regex = re.compile(regex_text)
+            r = re.match(regex, key)
+            if not r:
+                return False
+
+            m.clear()
+            m.extend(r.groups())
+            return True
+
+        state = {k:v for k, v in state.items() if k.startswith(prefix)}
+        sd_covert = {}
+        prefix_len = len(prefix)
+        for k, v in state.items():
+            model_k, lora_k = k[prefix_len:].split('.', 1)
+
+            model_k = self.replace_all(model_k, com_name, com_name_tmp).replace('_', '.').replace('%', '_')
+
+            if match(model_k, r'input_blocks.(\d+).1.(.+)'):
+                new_k = f'down_blocks.{down[m[0]][0]}.attentions.{down[m[0]][1]}.{m[1]}'
+            elif match(model_k, r'middle_block.1.(.+)'):
+                new_k = f'mid_block.attentions.0.{m[0]}'
+                pass
+            elif match(model_k, r'output_blocks.(\d+).(\d+).(.+)'):
+                new_k = f'up_blocks.{up[m[0]][0]}.attentions.{up[m[0]][1]}.{m[2]}'
+            else:
+                raise NotImplementedError
+
+            if lora_k == 'alpha':
+                sd_covert[f'{new_k}.___.{lora_k}'] = v
+            else:
+                sd_covert[f'{new_k}.___.layer.{self.lora_w_map[lora_k]}'] = v
+
+        return sd_covert
+
+    @staticmethod
+    def replace_all(data: str, srcs: List[str], dsts: List[str]):
+        for src, dst in zip(srcs, dsts):
+            data = data.replace(src, dst)
+        return data
+
+    @staticmethod
+    def alpha_scale_from_webui(state):
+        # Apply to "lora_down" and "lora_up" respectively to prevent overflow
+        for k, v in state.items():
+            if 'W_up' in k:
+                state[k] = v*math.sqrt(v.shape[1])
+            elif 'W_down' in k:
+                state[k] = v*math.sqrt(v.shape[0])
+        return state
+
+    @staticmethod
+    def alpha_scale_to_webui(state):
+        for k, v in state.items():
+            if 'lora_up' in k:
+                state[k] = v*math.sqrt(v.shape[1])
+            elif 'lora_down' in k:
+                state[k] = v*math.sqrt(v.shape[0])
+        return state
+
+class LoraWebuiFormat(CkptFormat):
+    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):
+        sd_denoiser = {k.removeprefix('denoiser.'):v for k, v in sd_model['base'].items() if k.startswith('denoiser.')}
+        sd_TE = {k.removeprefix('TE.'):v for k, v in sd_model['base'].items() if k.startswith('TE.')}
+
+        if len(sd_denoiser)>0 or len(sd_TE)>0:
+            sdxl = False
+            for k in sd_TE.keys():
+                if 'clip_L' in k or 'clip_bigG' in k:
+                    sdxl = True
+                    break
+            sd_webui = self.converter.convert_to_webui(sd_denoiser, sd_TE, auto_scale_alpha=self.auto_scale_alpha, sdxl=sdxl)
+        else:
+            sd_webui = self.converter.convert_to_webui(sd_model['base'], {}, auto_scale_alpha=self.auto_scale_alpha)
+
+        self.format.save_ckpt(sd_webui, save_f)
+
+    def load_ckpt(self, ckpt_f: str, map_location="cpu", **kwargs):
+        sd_webui = self.format.load_ckpt(ckpt_f, map_location=map_location, **kwargs)
+
+        sdxl = False
+        for k in sd_webui.keys():
+            if ('lora_te1_' in k or 'lora_te2_' in k or
+                    re.match(r'input_blocks.(\d+).1.(.+)', k) or
+                    re.match(r'middle_block.1.(.+)', k) or
+                    re.match(r'output_blocks.(\d+).(\d+).(.+)', k)):
+                sdxl = True
+                break
+
+        sd_all = self.converter.convert_from_webui(sd_webui, auto_scale_alpha=self.auto_scale_alpha, sdxl=sdxl)
+        return sd_all

hcpdiff/ckpt_manager/format/sd_single.py

@@ -0,0 +1,41 @@
+import torch
+from diffusers import AutoencoderKL, StableDiffusionPipeline, StableDiffusionXLPipeline
+from rainbowneko.ckpt_manager.format import CkptFormat
+
+from hcpdiff.diffusion.sampler import DDPMSampler, DDPMDiscreteSigmaScheduler
+from hcpdiff.models.compose import SDXLTextEncoder, SDXLTokenizer
+
+class OfficialSD15Format(CkptFormat):
+    # Single file format
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        pipe_args = dict(unet=denoiser, vae=vae, text_encoder=TE, tokenizer=tokenizer)
+        pipe_args = {k:v for k,v in pipe_args.items() if v is not None}
+        pipe = StableDiffusionPipeline.from_single_file(
+            pretrained_model, revision=revision, torch_dtype=dtype, **pipe_args
+        )
+        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        return dict(denoiser=pipe.unet, TE=pipe.text_encoder, vae=pipe.vae, noise_sampler=noise_sampler, tokenizer=pipe.tokenizer)
+
+class OfficialSDXLFormat(CkptFormat):
+    # Single file format
+    def load_ckpt(self, pretrained_model: str, map_location="cpu", denoiser=None, TE=None, vae: AutoencoderKL = None, noise_sampler=None,
+                  tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+        pipe_args = dict(unet=denoiser, vae=vae)
+        if TE is not None:
+            pipe_args['text_encoder'] = TE.clip_L
+            pipe_args['text_encoder_2'] = TE.clip_bigG
+        if tokenizer is not None:
+            pipe_args['tokenizer'] = tokenizer.clip_L
+            pipe_args['tokenizer_2'] = tokenizer.clip_bigG
+
+        pipe_args = {k:v for k,v in pipe_args.items() if v is not None}
+        pipe = StableDiffusionXLPipeline.from_single_file(
+            pretrained_model, revision=revision, torch_dtype=dtype, **pipe_args
+        )
+
+        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        TE = SDXLTextEncoder([('clip_L', pipe.text_encoder), ('clip_bigG', pipe.text_encoder_2)])
+        tokenizer = SDXLTokenizer([('clip_L', pipe.tokenizer), ('clip_bigG', pipe.tokenizer_2)])
+
+        return dict(denoiser=pipe.unet, TE=TE, vae=pipe.vae, noise_sampler=noise_sampler, tokenizer=tokenizer)

hcpdiff/ckpt_manager/loader.py

@@ -0,0 +1,64 @@
+from hcpdiff.models.lora_layers_patch import LoraLayer
+from torch import nn
+from hcpdiff.utils.net_utils import split_module_name
+from rainbowneko.ckpt_manager import NekoPluginLoader, LocalCkptSource, CkptFormat
+from rainbowneko.ckpt_manager.locator import get_match_layers
+from rainbowneko.models.plugin import PluginGroup
+
+def get_lora_rank_and_cls(lora_state):
+    if 'layer.W_down' in lora_state:
+        rank = lora_state['layer.W_down'].shape[0]
+        return LoraLayer, rank
+    else:
+        raise ValueError('Unknown lora format.')
+
+class HCPLoraLoader(NekoPluginLoader):
+    def __init__(self, format: CkptFormat=None, source: LocalCkptSource=None, path: str = None, layers='all', target_plugin=None,
+                 state_prefix=None, base_model_alpha=0.0, load_ema=False, module_to_load='', **plugin_kwargs):
+        super().__init__(format, source, path=path, layers=layers, target_plugin=target_plugin, state_prefix=state_prefix,
+                         base_model_alpha=base_model_alpha, load_ema=load_ema, **plugin_kwargs)
+        self.module_to_load = module_to_load
+
+    def load_to(self, name, model):
+        # get model to load plugin and its named_modules
+        model = model if self.module_to_load == '' else eval(f"model.{self.module_to_load}")
+
+        named_modules = {k:v for k, v in model.named_modules()}
+        plugin_state = self.load(self.path, map_location='cpu')['base_ema' if self.load_ema else 'base']
+
+        # filter layers to load
+        if self.layers != 'all':
+            match_blocks = get_match_layers(self.layers, named_modules)
+            plugin_state = {k: v for blk in match_blocks for k, v in plugin_state.items() if k.startswith(blk)}
+
+        if self.state_prefix:
+            state_prefix_len = len(self.state_prefix)
+            plugin_state = {k[state_prefix_len:]: v for k, v in plugin_state.items() if k.startswith(self.state_prefix)}
+
+        lora_block_state = {}
+        # get all layers in the lora_state
+        for pname, p in plugin_state.items():
+            # lora_block. is the old format
+            prefix, block_name = pname.split('.___.', 1)
+            if prefix not in lora_block_state:
+                lora_block_state[prefix] = {}
+            lora_block_state[prefix][block_name] = p
+
+        # add lora to host and load weights
+        lora_blocks = {}
+        for layer_name, lora_state in lora_block_state.items():
+            lora_layer_cls, rank = get_lora_rank_and_cls(lora_state)
+
+            if 'alpha' in lora_state:
+                lora_state['alpha'] *= self.plugin_kwargs.get('alpha', 1.0)
+
+            parent_name, host_name = split_module_name(layer_name)
+
+            lora_block = lora_layer_cls.wrap_layer(name, named_modules[layer_name], rank=rank, bias='layer.bias' in lora_state,
+                                                parent_block=named_modules[parent_name], host_name=host_name)
+            lora_block.set_hyper_params(**self.plugin_kwargs)
+            lora_blocks[layer_name] = lora_block
+            load_info = lora_block.load_state_dict(lora_state, strict=False)
+            if len(load_info.unexpected_keys) > 0:
+                print(name, 'unexpected_keys', load_info.unexpected_keys)
+        return PluginGroup(lora_blocks)

hcpdiff/data/__init__.py

@@ -1,28 +1,4 @@
-from .pair_dataset import TextImagePairDataset
-from .cond_dataset import TextImageCondPairDataset
-from .crop_info_dataset import CropInfoPairDataset
-from .bucket import BaseBucket, FixedBucket, RatioBucket, SizeBucket, RatioSizeBucket, LongEdgeBucket
-from .utils import CycleData
-from .caption_loader import JsonCaptionLoader, TXTCaptionLoader
-from .sampler import DistributedCycleSampler, get_sampler
-
-class DataGroup:
-    def __init__(self, loader_list, loss_weights):
-        self.loader_list = loader_list
-        self.loss_weights = loss_weights
-
-    def __iter__(self):
-        self.data_iter_list = [iter(CycleData(loader)) for loader in self.loader_list]
-        return self
-
-    def __next__(self):
-        return [next(data_iter) for data_iter in self.data_iter_list]
-
-    def __len__(self):
-        return len(self.loader_list)
-
-    def get_dataset(self, idx):
-        return self.loader_list[idx].dataset
-
-    def get_loss_weights(self, idx):
-        return self.loss_weights[idx]
+from .dataset import TextImagePairDataset
+from .source import Text2ImageSource, Text2ImageLossMapSource, Text2ImageCondSource, T2IFolderClassSource, TextSource
+from .handler import StableDiffusionHandler, LossMapHandler, DiffusionImageHandler, DiffusionTextHandler
+from .cache import VaeCache

hcpdiff/data/cache/__init__.py

@@ -0,0 +1 @@
+from .vae import VaeCache

hcpdiff/data/cache/vae.py

@@ -0,0 +1,102 @@
+from io import BytesIO
+from pathlib import Path
+from typing import Dict, Any
+
+import lmdb
+import torch
+from hcpdiff.models.wrapper import SD15Wrapper
+from rainbowneko import _share
+from rainbowneko.data import DataCache, CacheableDataset
+from rainbowneko.utils import Path_Like
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+class VaeCache(DataCache):
+    def __init__(self, pre_build: Path_Like = None, lazy=False, bs=1):
+        super().__init__(pre_build)
+        self.lazy = lazy
+        self.bs = bs
+
+    def load_latent(self, id):
+        if self.lazy:
+            with self.env.begin() as txn:
+                byte_tensor = txn.get(str(id).encode())
+                return torch.load(BytesIO(byte_tensor))
+        else:
+            return self.cache[id]
+
+    def before_handler(self, index: int, data: Dict[str, Any]):
+        cached_data = self.load_latent(data['id'])
+        data['image'] = cached_data['latent']
+        data['coord'] = cached_data['coord']
+        return data
+    
+    def on_finish(self, index, data):
+        return data
+
+    def load(self, path):
+        if self.lazy:
+            self.env = lmdb.open(path, subdir=False, readonly=True, lock=False, readahead=False, meminit=False)
+            return {}
+        elif len(self.cache)>0:
+            return self.cache
+        else:
+            env = lmdb.open(path, subdir=False, readonly=True, lock=False, readahead=False, meminit=False)
+            with env.begin() as txn:
+                cache = {k.decode():torch.load(BytesIO(v)) for k, v in txn.cursor()}
+            env.close()
+            return cache
+
+    def build(self, dataset: CacheableDataset, model: SD15Wrapper, all_gather):
+        if (self.pre_build and Path(self.pre_build).exists()) or len(self.cache)>0:
+            model.vae = None
+            return
+
+        vae = model.vae.to(_share.device)
+        with dataset.disable_cache():
+            dataset.bucket.rest(0)
+
+            loader = DataLoader(
+                dataset,
+                batch_size=self.bs,
+                num_workers=0,
+                sampler=DistributedSampler(dataset, num_replicas=_share.world_size, rank=_share.local_rank, shuffle=False),
+                collate_fn=dataset.collate_fn,
+                drop_last=False,
+            )
+
+            if self.pre_build:
+                Path(self.pre_build).parent.mkdir(parents=True, exist_ok=True)
+                env = lmdb.open(self.pre_build, map_size=1099511627776)
+                with env.begin(write=True) as txn:
+                    for data in tqdm(loader):
+                        image = data['image'].to(device=_share.device, dtype=vae.dtype)
+                        latents = model.vae.encode(image).latent_dist.sample()
+                        latents = (latents*vae.config.scaling_factor).cpu()
+
+                        for img_id, latent, coord in zip(data['id'], latents, data['coord']):
+                            data_cache = {'latent': latent, 'coord': coord}
+
+                            byte_stream = BytesIO()
+                            torch.save(data_cache, byte_stream)
+                            txn.put(str(img_id).encode(), byte_stream.getvalue())
+                            if not self.lazy:
+                                self.cache[img_id] = data_cache
+                env.close()
+            else:
+                for data in tqdm(loader):
+                    img_id = data['id']
+                    image = data['image'].to(device=_share.device, dtype=vae.dtype)
+                    latents = model.vae.encode(image).latent_dist.sample()
+                    latents = (latents*vae.config.scaling_factor).cpu()
+                    for img_id, latent, coord in zip(data['id'], latents, data['coord']):
+                        self.cache[img_id] = {'latent': latent, 'coord': coord}
+
+        model.vae.to('cpu')
+        #model.vae = None
+        torch.cuda.empty_cache()
+
+        cache_all = all_gather(self.cache)
+        for cache in cache_all:
+            self.cache.update(cache)

hcpdiff/data/dataset.py

@@ -0,0 +1,20 @@
+"""
+pair_dataset.py
+====================
+    :Name:        text-image pair dataset
+    :Author:      Dong Ziyi
+    :Affiliation: HCP Lab, SYSU
+    :Created:     10/03/2023
+    :Licence:     Apache-2.0
+"""
+
+from typing import Union, Dict
+
+from rainbowneko.data import CacheableDataset, BaseDataset, BaseBucket, DataSource, DataHandler, DataCache
+
+def TextImagePairDataset(bucket: BaseBucket = None, source: Dict[str, DataSource] = None, handler: DataHandler = None,
+                         batch_handler: DataHandler = None, cache: DataCache = None, **kwargs) -> Union[CacheableDataset, BaseDataset]:
+    if cache is None:
+        return BaseDataset(bucket=bucket, source=source, handler=handler, batch_handler=batch_handler, **kwargs)
+    else:
+        return CacheableDataset(bucket=bucket, source=source, handler=handler, batch_handler=batch_handler, cache=cache, **kwargs)

hcpdiff/data/handler/__init__.py

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