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

hcpdiff/easy/cfg/sdxl_train.py

@@ -0,0 +1,147 @@
+import torch
+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,
+                    dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0, name: str = 'SDXL'):
+    if low_vram:
+        from bitsandbytes.optim import AdamW8bit
+        optimizer = AdamW8bit(_partial_=True)
+    else:
+        optimizer = torch.optim.AdamW(_partial_=True)
+
+    from cfgs.train.py import train_base, tuning_base
+
+    return dict(
+        _base_=[train_base, tuning_base],
+        mixed_precision=dtype,
+
+        model_part=CfgWDModelParser([
+            dict(
+                lr=lr,
+                layers=['denoiser'],  # train UNet
+            )
+        ], weight_decay=1e-2),
+
+        ckpt_saver=dict(
+            SDXL=ckpt_saver(
+                ckpt_type='safetensors',
+                target_module='denoiser',
+                layers=LAYERS_TRAINABLE,
+            )
+        ),
+
+        train=dict(
+            train_steps=train_steps,
+            save_step=save_step,
+
+            optimizer=optimizer,
+
+            scheduler=ConstantLR(
+                _partial_=True,
+                warmup_steps=warmup_steps,
+            ),
+        ),
+
+        model=dict(
+            name=name,
+
+            ## Easy config
+            wrapper=SDXLWrapper.from_pretrained(
+                _partial_=True,
+                models=SDXL_auto_loader(ckpt_path=base_model, _partial_=True),
+            ),
+        ),
+
+        data_train=dataset,
+    )
+
+@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 = 'SDXL',
+                    save_webui_format=False):
+    with disable_neko_cfg:
+        if alpha is None:
+            alpha = rank
+
+        if with_conv:
+            lora_layers = [
+                r're:denoiser.*\.attn.?$',
+                r're:denoiser.*\.ff$',
+                r're:denoiser.*\.resnets$',
+                r're:denoiser.*\.proj_in$',
+                r're:denoiser.*\.proj_out$',
+                r're:denoiser.*\.conv$',
+            ]
+        else:
+            lora_layers = [
+                r're:denoiser.*\.attn.?$',
+                r're:denoiser.*\.ff$',
+            ]
+
+    if low_vram:
+        from bitsandbytes.optim import AdamW8bit
+        optimizer = AdamW8bit(_partial_=True, betas=(0.9, 0.99))
+    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(
+        _base_=[SD_FT],
+        mixed_precision=dtype,
+
+        model_part=None,
+        model_plugin=CfgWDPluginParser(cfg_plugin=dict(
+            lora1=LoraLayer.wrap_model(
+                _partial_=True,
+                lr=lr,
+                rank=rank,
+                alpha=alpha,
+                layers=lora_layers
+            )
+        ), weight_decay=0.1),
+
+        ckpt_saver=dict(
+            _replace_ = True,
+            lora_unet=NekoPluginSaver(
+                format=lora_format,
+                target_plugin='lora1',
+            )
+        ),
+
+        train=dict(
+            train_steps=train_steps,
+            save_step=save_step,
+
+            optimizer=optimizer,
+
+            scheduler=ConstantLR(
+                _partial_=True,
+                warmup_steps=warmup_steps,
+            ),
+        ),
+
+        model=dict(
+            name=name,
+
+            wrapper=SDXLWrapper.from_pretrained(
+                models=SDXL_auto_loader(ckpt_path=base_model, _partial_=True),
+                _partial_=True,
+            ),
+        ),
+
+        data_train=dataset,
+    )

hcpdiff/easy/cfg/t2i.py

@@ -0,0 +1,228 @@
+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
+
+from hcpdiff.ckpt_manager import HCPLoraLoader
+from hcpdiff.easy import Diffusers_SD, SD15_auto_loader, SDXL_auto_loader
+from hcpdiff.workflow import (BuildModelsAction, PrepareDiffusionAction, XformersEnableAction, VaeOptimizeAction, TextHookAction,
+                              AttnMultTextEncodeAction, SeedAction, MakeTimestepsAction, MakeLatentAction, DiffusionStepAction,
+                              time_iter, DecodeAction, SaveImageAction, LatentResizeAction)
+
+negative_prompt = 'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'
+
+## Easy config
+@neko_cfg
+def build_model(pretrained_model='ckpts/any5', noise_sampler=Diffusers_SD.dpmpp_2m_karras) -> Actions:
+    return Actions([
+        PrepareAction(device='cuda', dtype=torch.float16),
+        BuildModelsAction(
+            model_loader=SD15_auto_loader(
+                _partial_=True,
+                ckpt_path=pretrained_model,
+                noise_sampler=noise_sampler
+            )
+        ),
+    ])
+
+@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 = []
+    for i, item in enumerate(info):
+        lora_unet = LoadModelAction(cfg={
+            f'lora_unet_{i}':HCPLoraLoader(
+                path=item[0],
+                state_prefix='denoiser.',
+                alpha=item[1],
+            )
+        }, key_map_in=('denoiser -> model', 'in_preview -> in_preview'))
+        lora_TE = LoadModelAction(cfg={
+            f'lora_TE_{i}':HCPLoraLoader(
+                path=item[0],
+                state_prefix='TE.',
+                alpha=item[1],
+            )
+        }, key_map_in=('TE -> model', 'in_preview -> in_preview'))
+
+        with disable_neko_cfg:
+            lora_acts.append(lora_unet)
+            lora_acts.append(lora_TE)
+
+    return Actions(lora_acts)
+
+@neko_cfg
+def optimize_model() -> Actions:
+    return Actions([
+        PrepareDiffusionAction(),
+        XformersEnableAction(),
+        VaeOptimizeAction(slicing=True),
+    ])
+
+@neko_cfg
+def text(prompt, negative_prompt=negative_prompt, bs=4, N_repeats=1, layer_skip=1) -> Actions:
+    return Actions([
+        TextHookAction(N_repeats=N_repeats, layer_skip=layer_skip),
+        AttnMultTextEncodeAction(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            bs=bs
+        ),
+    ])
+
+@neko_cfg
+def build_model_SDXL(pretrained_model='ckpts/any5', noise_sampler=Diffusers_SD.dpmpp_2m_karras) -> Actions:
+    return Actions([
+        PrepareAction(device='cuda', dtype=torch.float16),
+        ## Easy config
+        BuildModelsAction(
+            model_loader=SDXL_auto_loader(
+                _partial_=True,
+                ckpt_path=pretrained_model,
+                noise_sampler=noise_sampler
+            )
+        ),
+    ])
+
+@neko_cfg
+def text_SDXL(prompt, negative_prompt=negative_prompt, bs=4, N_repeats=1, layer_skip=1) -> Actions:
+    return Actions([
+        TextHookAction(N_repeats=N_repeats, layer_skip=layer_skip, TE_final_norm=False),
+        AttnMultTextEncodeAction(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            bs=bs
+        ),
+    ])
+
+@neko_cfg
+def config_diffusion(width=512, height=512, seed=None, N_steps=20, strength: float = None) -> Actions:
+    return Actions([
+        SeedAction(seed),
+        MakeTimestepsAction(N_steps=N_steps, strength=strength),
+        MakeLatentAction(width=width, height=height)
+    ])
+
+@neko_cfg
+def diffusion(guidance_scale=7.0) -> Actions:
+    return Actions([
+        LoopAction(
+            iterator=time_iter,
+            actions=[
+                DiffusionStepAction(guidance_scale=guidance_scale)
+            ]
+        )
+    ])
+
+@neko_cfg
+def decode(save_root='output_pipe/') -> Actions:
+    return Actions([
+        DecodeAction(),
+        SaveImageAction(save_root=save_root, image_type='png'),
+    ])
+
+@neko_cfg
+def resize(width=1024, height=1024):
+    return Actions([
+        LatentResizeAction(width=width, height=height)
+    ])
+
+@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/', 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, 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_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/', 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, 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(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/', 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, 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/', 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, 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/easy/model/__init__.py

@@ -0,0 +1,2 @@
+from .loader import SD15_auto_loader, SDXL_auto_loader, PixArt_auto_loader
+from .cnet import ControlNet_SD15, make_controlnet_handler

hcpdiff/easy/model/cnet.py

@@ -0,0 +1,31 @@
+from hcpdiff.data.handler import ControlNetHandler, StableDiffusionHandler
+from hcpdiff.models import ControlNetPlugin
+from rainbowneko.data import SyncHandler
+from rainbowneko.parser import neko_cfg
+
+@neko_cfg
+def ControlNet_SD15(lr=1e-4):
+    return ControlNetPlugin(
+        _partial_=True,
+        lr=lr,
+        from_layers=[
+            'pre_hook:',
+            'pre_hook:conv_in',  # to make forward inside autocast
+        ],
+        to_layers=[
+            'down_blocks.0',
+            'down_blocks.1',
+            'down_blocks.2',
+            'down_blocks.3',
+            'mid_block',
+            'pre_hook:up_blocks.3.resnets.2',
+        ]
+    )
+
+@neko_cfg
+def make_controlnet_handler(bucket=None, encoder_attention_mask=False, erase=0.15, dropout=0.0, shuffle=0.0, word_names={}):
+    return SyncHandler(
+        diffusion=StableDiffusionHandler(bucket=bucket, encoder_attention_mask=encoder_attention_mask, erase=erase, dropout=dropout, shuffle=shuffle,
+                                         word_names=word_names),
+        cnet=ControlNetHandler(bucket=bucket)
+    )

hcpdiff/easy/model/loader.py

@@ -0,0 +1,79 @@
+import torch
+from hcpdiff.ckpt_manager import DiffusersSD15Format, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSD15Format, OfficialSDXLFormat
+from rainbowneko.ckpt_manager import NekoLoader, LocalCkptSource
+from hcpdiff.utils import auto_tokenizer_cls, auto_text_encoder_cls, get_pipe_name
+from hcpdiff.models.wrapper import SDXLWrapper, SD15Wrapper, PixArtWrapper
+from hcpdiff.models.compose import SDXLTextEncoder
+from diffusers import  StableDiffusionPipeline, StableDiffusionXLPipeline
+
+def SD15_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    try:
+        try_diffusers = StableDiffusionPipeline.load_config(ckpt_path)
+        loader = NekoLoader(
+            format=DiffusersSD15Format(),
+            source=LocalCkptSource(),
+        )
+    except EnvironmentError:
+        loader = NekoLoader(
+            format=OfficialSD15Format(),
+            source=LocalCkptSource(),
+        )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype, **kwargs)
+    return models
+
+def SDXL_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    try:
+        try_diffusers = StableDiffusionXLPipeline.load_config(ckpt_path)
+        loader = NekoLoader(
+            format=DiffusersSDXLFormat(),
+            source=LocalCkptSource(),
+        )
+    except EnvironmentError:
+        loader = NekoLoader(
+            format=OfficialSDXLFormat(),
+            source=LocalCkptSource(),
+        )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype, **kwargs)
+    return models
+
+def PixArt_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    loader = NekoLoader(
+        format=DiffusersPixArtFormat(),
+        source=LocalCkptSource(),
+    )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype, **kwargs)
+    return models
+
+def auto_load_wrapper(pretrained_model, denoiser=None, TE=None, vae=None, noise_sampler=None, tokenizer=None, revision=None,
+                      dtype=torch.float32, **kwargs):
+    if TE is not None:
+        text_encoder_cls = type(TE)
+    else:
+        text_encoder_cls = auto_text_encoder_cls(pretrained_model, revision)
+
+    pipe_name = get_pipe_name(pretrained_model)
+
+    if text_encoder_cls == SDXLTextEncoder:
+        wrapper_cls = SDXLWrapper
+        format = DiffusersSDXLFormat()
+    elif 'PixArt' in pipe_name:
+        wrapper_cls = PixArtWrapper
+        format = DiffusersPixArtFormat()
+    else:
+        wrapper_cls = SD15Wrapper
+        format = DiffusersSD15Format()
+
+    loader = NekoLoader(
+        format=format,
+        source=LocalCkptSource(),
+    )
+    models = loader.load(pretrained_model, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                          dtype=dtype)
+
+    return wrapper_cls.build_from_pretrained(models, **kwargs)

hcpdiff/easy/sampler.py

@@ -0,0 +1,46 @@
+from hcpdiff.diffusion.sampler import DiffusersSampler
+from diffusers import DPMSolverMultistepScheduler, DDIMScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler
+
+class Diffusers_SD:
+    dpmpp_2m = DiffusersSampler(
+        DPMSolverMultistepScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+            algorithm_type='dpmsolver++',
+        )
+    )
+
+    dpmpp_2m_karras = DiffusersSampler(
+        DPMSolverMultistepScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+            algorithm_type='dpmsolver++',
+            use_karras_sigmas=True,
+        )
+    )
+
+    ddim = DiffusersSampler(
+        DDIMScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+        )
+    )
+
+    euler = DiffusersSampler(
+        EulerDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+        )
+    )
+
+    euler_a = DiffusersSampler(
+        EulerAncestralDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule='scaled_linear',
+        )
+    )

hcpdiff/evaluate/__init__.py

@@ -0,0 +1 @@
+from .previewer import HCPPreviewer

hcpdiff/evaluate/previewer.py

@@ -0,0 +1,60 @@
+from pathlib import Path
+
+import torch
+from rainbowneko.evaluate.preview import WorkflowPreviewer
+from rainbowneko.utils import to_cuda
+
+from hcpdiff.models.wrapper import SD15Wrapper
+from accelerate.hooks import remove_hook_from_module
+
+class HCPPreviewer(WorkflowPreviewer):
+
+    @torch.no_grad()
+    def evaluate(self, step: int, model: SD15Wrapper, prefix='eval/'):
+        if step%self.interval != 0 or not self.trainer.is_local_main_process:
+            return
+
+        # record training layers
+        training_layers = [layer for layer in model.modules() if layer.training]
+
+        model.eval()
+        self.trainer.loggers.info(f'Preview')
+
+        N_repeats = model.text_enc_hook.N_repeats
+        clip_skip = model.text_enc_hook.clip_skip
+        clip_final_norm = model.text_enc_hook.clip_final_norm
+        use_attention_mask = model.text_enc_hook.use_attention_mask
+
+        preview_root = Path(self.trainer.exp_dir)/'imgs'
+        preview_root.mkdir(parents=True, exist_ok=True)
+
+        states = self.workflow_runner.run(model=model, in_preview=True, te_hook=model.text_enc_hook,
+                                          device=self.device, dtype=self.dtype, preview_root=preview_root, preview_step=step,
+                                          world_size=self.trainer.world_size, local_rank=self.trainer.local_rank,
+                                          emb_hook=self.trainer.cfgs.emb_pt.embedding_hook if self.trainer.pt_trainable else None)
+
+        # restore model states
+        if model.vae is not None:
+            model.vae.disable_tiling()
+            model.vae.disable_slicing()
+            remove_hook_from_module(model.vae, recurse=True)
+            if 'vae_encode_raw' in states:
+                model.vae.encode = states['vae_encode_raw']
+                model.vae.decode = states['vae_decode_raw']
+
+        if 'emb_hook' in states and not self.trainer.pt_trainable:
+            states['emb_hook'].remove()
+
+        if self.trainer.pt_trainable:
+            self.trainer.cfgs.emb_pt.embedding_hook.N_repeats = N_repeats
+
+        model.tokenizer.N_repeats = N_repeats
+        model.text_enc_hook.N_repeats = N_repeats
+        model.text_enc_hook.clip_skip = clip_skip
+        model.text_enc_hook.clip_final_norm = clip_final_norm
+        model.text_enc_hook.use_attention_mask = use_attention_mask
+        
+        to_cuda(model)
+
+        for layer in training_layers:
+            layer.train()

hcpdiff/loss/__init__.py

@@ -1 +1,4 @@
-from .min_snr_loss import MinSNRLoss, SoftMinSNRLoss, KDiffMinSNRLoss, EDMLoss
+from .weighting import MinSNRWeight, SNRWeight, EDMWeight, LossWeight
+from .ssim import SSIMLoss, MS_SSIMLoss
+from .gw import GWLoss
+from .base import DiffusionLossContainer

hcpdiff/loss/base.py

@@ -0,0 +1,41 @@
+from rainbowneko.train.loss import LossContainer
+from typing import Dict, Any
+from torch import Tensor
+
+class DiffusionLossContainer(LossContainer):
+    def __init__(self, loss, weight=1.0, key_map=None):
+        key_map = key_map or getattr(loss, '_key_map', None) or ('pred.model_pred -> 0', 'pred.target -> 1')
+        super().__init__(loss, weight, key_map)
+        self.target_type = getattr(loss, 'target_type', 'eps')
+
+    def get_target(self, pred_type, model_pred, x_0, noise, x_t, sigma, noise_sampler, **kwargs):
+        # Get target
+        if self.target_type == "eps":
+            target = noise
+        elif self.target_type == "x0":
+            target = x_0
+        elif self.target_type == "velocity":
+            target = noise_sampler.eps_to_velocity(noise, x_t, sigma)
+        else:
+            raise ValueError(f"Unsupport target_type {self.target_type}")
+
+        # TODO: put in wrapper
+        # # remove pred vars
+        # if model_pred.shape[1] == target.shape[1]*2:
+        #     model_pred, _ = model_pred.chunk(2, dim=1)
+
+        # Convert pred_type to target_type
+        if pred_type != self.target_type:
+            cvt_func = getattr(noise_sampler, f'{pred_type}_to_{self.target_type}', None)
+            if cvt_func is None:
+                raise ValueError(f"Unsupport pred_type {pred_type} with target_type {self.target_type}")
+            else:
+                model_pred = cvt_func(model_pred, x_t, sigma)
+        return model_pred, target
+    
+    def forward(self, pred:Dict[str,Any], inputs:Dict[str,Any]) -> Tensor:
+        model_pred, target = self.get_target(**pred)
+        pred['model_pred'] = model_pred
+        pred['target'] = target
+        loss = super().forward(pred, inputs) * self.weight # [B,*,*,*]
+        return loss.mean()

hcpdiff/loss/gw.py

@@ -0,0 +1,35 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class GWLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        sobel_x = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
+        sobel_y = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]
+        self.sobel_x = torch.FloatTensor(sobel_x)
+        self.sobel_y = torch.FloatTensor(sobel_y)
+        self.register_buffer('sobel_x', self.sobel_x)
+        self.register_buffer('sobel_y', self.sobel_y)
+
+    def forward(self, pred, target):
+        '''
+
+        :param pred: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,C,H,W]
+        '''
+        b, c, w, h = pred.shape
+
+        sobel_x = self.sobel_x.expand(c, 1, 3, 3).to(pred.device)
+        sobel_y = self.sobel_y.expand(c, 1, 3, 3).to(pred.device)
+        Ix1 = F.conv2d(pred, sobel_x, stride=1, padding=1, groups=c)
+        Ix2 = F.conv2d(target, sobel_x, stride=1, padding=1, groups=c)
+        Iy1 = F.conv2d(pred, sobel_y, stride=1, padding=1, groups=c)
+        Iy2 = F.conv2d(target, sobel_y, stride=1, padding=1, groups=c)
+
+        dx = torch.abs(Ix1 - Ix2)
+        dy = torch.abs(Iy1 - Iy2)
+        loss = (1 + 4 * dx) * (1 + 4 * dy) * torch.abs(pred - target)
+        return loss