hcpdiff 0.3.5__tar.gz → 0.3.7__tar.gz

{hcpdiff-0.3.5/hcpdiff.egg-info → hcpdiff-0.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hcpdiff
-Version: 0.3.5
+Version: 0.3.7
 Summary: A universal Stable-Diffusion toolbox
 Home-page: https://github.com/7eu7d7/HCP-Diffusion
 Author: Ziyi Dong
@@ -58,6 +58,8 @@ Compared to DreamArtist, DreamArtist++ is more stable with higher image quality
 * safetensors support
 * Controlnet (support train)
 * Min-SNR loss
+* Custom optimizer (Lion, DAdaptation, pytorch-optimizer, ...)
+* Custom lr scheduler
 
 ## Install
 

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/README.md

@@ -39,6 +39,8 @@ Compared to DreamArtist, DreamArtist++ is more stable with higher image quality
 * safetensors support
 * Controlnet (support train)
 * Min-SNR loss
+* Custom optimizer (Lion, DAdaptation, pytorch-optimizer, ...)
+* Custom lr scheduler
 
 ## Install
 

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/cfgs/train/examples/DreamArtist++.yaml

@@ -6,15 +6,21 @@ lora_unet:
   - lr: 1e-4
     rank: 0.01875
     branch: p
+    dropout: 0.1
     layers:
-      - 're:.*\.attn.?$'
-      #- 're:.*\.ff\.net\.0$' # Increases fitness, but potentially reduces controllability
+      - 're:.*\.to_k$'
+      - 're:.*\.to_v$'
+      - 're:.*\.ff$'
+      #- 're:.*\.attn.?$' # Increases fitness, but potentially reduces controllability
   - lr: 4e-5 # Low negative unet lr prevents image collapse
     rank: 0.01875
     branch: n
+    dropout: 0.1
     layers:
-      - 're:.*\.attn.?$'
-      #- 're:.*\.ff\.net\.0$' # Increases fitness, but potentially reduces controllability
+      - 're:.*\.to_k$'
+      - 're:.*\.to_v$'
+      - 're:.*\.ff$'
+      #- 're:.*\.attn.?$' # Increases fitness, but potentially reduces controllability
   #  - lr: 1e-4
   #    rank: 0.01875
   #    type: p
@@ -27,23 +33,25 @@ lora_unet:
   #      - 're:.*\.resnets$' # Increases fitness, but potentially reduces controllability and change style
 
 lora_text_encoder:
-  - lr: 1e-5
-    rank: 0.01
+  - lr: 2e-5
+    rank: 2
     branch: p
+    dropout: 0.1
     layers:
       - 're:.*self_attn$'
       - 're:.*mlp$'
-  - lr: 1e-5
-    rank: 0.01
+  - lr: 2e-5
+    rank: 2
     branch: n
+    dropout: 0.1
     layers:
       - 're:.*self_attn$'
       - 're:.*mlp$'
 
 tokenizer_pt:
   train: # prompt tuning embeddings
-    - { name: 'pt-botdog1', lr: 0.003 }
-    - { name: 'pt-botdog1-neg', lr: 0.003 }
+    - { name: 'pt-botdog1', lr: 0.0025 }
+    - { name: 'pt-botdog1-neg', lr: 0.0025 }
 
 train:
   gradient_accumulation_steps: 1
@@ -52,13 +60,19 @@ train:
   #cfg_scale: '1.0-3.0:cos' # dynamic CFG with timestamp
   cfg_scale: '3.0'
 
+  loss:
+    criterion: # min SNR loss
+      _target_: hcpdiff.loss.MinSNRLoss
+      gamma: 2.0
+
   scheduler:
-    name: 'constant_with_warmup'
-    num_warmup_steps: 50
+    name: one_cycle
+    num_warmup_steps: 200
     num_training_steps: 1000
+    scheduler_kwargs: { }
 
   scheduler_pt:
-    name: 'one_cycle'
+    name: one_cycle
     num_warmup_steps: 200
     num_training_steps: 1000
     scheduler_kwargs: {}
@@ -68,6 +82,7 @@ model:
   tokenizer_repeats: 1
   ema_unet: 0
   ema_text_encoder: 0
+  clip_skip: 0
 
 data:
   dataset1:

hcpdiff-0.3.7/cfgs/train/examples/Lion_optimizer.yaml

@@ -0,0 +1,15 @@
+_base_: [cfgs/train/examples/fine-tuning.yaml]
+
+# Install: pip install lion-pytorch
+
+train:
+  optimizer:
+    _target_: lion_pytorch.Lion
+    _partial_: True
+    weight_decay: 1e-2
+    #use_triton: True # set this to True to use cuda kernel w/ Triton lang (Tillet et al)
+
+  optimizer_pt:
+    _target_: lion_pytorch.Lion
+    _partial_: True
+    weight_decay: 1e-3

hcpdiff-0.3.7/cfgs/train/examples/min_snr.yaml

@@ -0,0 +1,7 @@
+_base_: [cfgs/train/examples/fine-tuning.yaml]
+
+train:
+  loss:
+    criterion: # min SNR loss
+      _target_: hcpdiff.loss.MinSNRLoss
+      gamma: 2.0

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/cfgs/train/train_base.yaml

@@ -10,6 +10,7 @@ vis_info:
   negative_prompt: ''
 
 train:
+  train_steps: 1000
   gradient_accumulation_steps: 1
   workers: 4
   max_grad_norm: 1.0
@@ -35,7 +36,10 @@ train:
   optimizer:
     type: adamw
     weight_decay: 1e-3
-    weight_decay_pt: 5e-4
+
+  optimizer_pt:
+    type: adamw
+    weight_decay: 5e-4
 
   scale_lr: True # auto scale lr with total batch size
   scheduler:
@@ -58,7 +62,7 @@ model:
   revision: null
   pretrained_model_name_or_path: null
   tokenizer_name: null
-  tokenizer_repeats: 3
+  tokenizer_repeats: 2
   enable_xformers: True
   gradient_checkpointing: True
   ema_unet: 0 # 0 to disable

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff/data/bucket.py

@@ -194,7 +194,7 @@ class RatioBucket(BaseBucket):
             rs.shuffle(x)
 
         # shuffle of batches
-        bucket_list = np.hstack(bucket_list).reshape(-1, self.bs)
+        bucket_list = np.hstack(bucket_list).reshape(-1, self.bs).astype(int)
         rs.shuffle(bucket_list)
 
         self.idx_arb = bucket_list.reshape(-1)

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff/data/utils.py

@@ -72,7 +72,6 @@ def collate_fn_ft(batch):
 
 class CycleData():
     def __init__(self, data_loader):
-        print(data_loader)
         self.data_loader = data_loader
 
     def __iter__(self):

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff/tools/convert_caption_txt2json.py

@@ -6,6 +6,7 @@ from hcpdiff.utils.img_size_tool import types_support
 
 parser = argparse.ArgumentParser(description='Stable Diffusion Training')
 parser.add_argument('--data_root', type=str, default='')
+parser.add_argument('--with_imgs', action="store_true")
 args = parser.parse_args()
 
 
@@ -16,10 +17,13 @@ def get_txt_caption(path):
 
 captions = {}
 for file in os.listdir(args.data_root):
-    ext_idx = file.rfind('.')
-    file_name = file[:ext_idx]
-    if file[ext_idx + 1:] in types_support:
-        captions[file] = get_txt_caption(os.path.join(args.data_root, f'{file_name}.txt'))
+    file_name, file_ext = file.rsplit('.', 1)
+    if args.with_imgs:
+        if file_ext in types_support:
+            captions[file] = get_txt_caption(os.path.join(args.data_root, f'{file_name}.txt'))
+    else:
+        if file_ext == 'txt':
+            captions[file] = get_txt_caption(os.path.join(args.data_root, f'{file_name}.txt'))
 
 with open(os.path.join(args.data_root, f'image_captions.json'), "w", encoding='utf8') as f:
-    json.dump(captions, f)
+    json.dump(captions, f, indent=2, ensure_ascii=False)

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff/train_ac.py

@@ -28,6 +28,7 @@ from diffusers import AutoencoderKL, UNet2DConditionModel
 from diffusers.utils.import_utils import is_xformers_available
 from omegaconf import OmegaConf
 from transformers import AutoTokenizer
+from  functools import partial
 
 from hcpdiff.ckpt_manager import CkptManagerPKL, CkptManagerSafe
 from hcpdiff.data import RatioBucket, DataGroup, collate_fn_ft
@@ -329,12 +330,16 @@ class Trainer:
         # set optimizer
         parameters, parameters_pt = self.get_param_group_train()
 
-        cfg_opt = self.cfgs.train.optimizer
         if len(parameters)>0:  # do fine-tuning
+            cfg_opt = self.cfgs.train.optimizer
             if self.cfgs.train.scale_lr:
                 self.scale_lr(parameters)
 
-            if cfg_opt.type == 'adamw_8bit':
+            if isinstance(cfg_opt, partial):
+                if 'type' in cfg_opt.keywords:
+                    del cfg_opt.keywords['type']
+                self.optimizer = cfg_opt(params=parameters, lr=self.lr)
+            elif cfg_opt.type == 'adamw_8bit':
                 import bitsandbytes as bnb
                 self.optimizer = bnb.optim.AdamW8bit(params=parameters, lr=self.lr, weight_decay=cfg_opt.weight_decay)
             elif cfg_opt.type == 'deepspeed' and self.accelerator.state.deepspeed_plugin is not None:
@@ -343,23 +348,35 @@ class Trainer:
             elif cfg_opt.type == 'adamw':
                 self.optimizer = torch.optim.AdamW(params=parameters, lr=self.lr, weight_decay=cfg_opt.weight_decay)
             else:
-                self.optimizer = cfg_opt.optimizer.opt(parameters, lr=self.lr)
+                raise NotImplementedError(f'Unknown optimizer {cfg_opt.type}')
 
-            self.lr_scheduler = get_scheduler(optimizer=self.optimizer, **self.cfgs.train.scheduler)
+            if isinstance(self.cfgs.train.scheduler, partial):
+                self.lr_scheduler = self.cfgs.train.scheduler(optimizer=self.optimizer)
+            else:
+                self.lr_scheduler = get_scheduler(optimizer=self.optimizer, **self.cfgs.train.scheduler)
 
         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)
+            if isinstance(cfg_opt_pt, partial):
+                if 'type' in cfg_opt_pt.keywords:
+                    del cfg_opt_pt.keywords['type']
+                self.optimizer_pt = cfg_opt_pt(params=parameters_pt, lr=self.lr)
+            else:
+                self.optimizer_pt = torch.optim.AdamW(params=parameters_pt, lr=self.lr, weight_decay=cfg_opt_pt.weight_decay)
 
-            self.optimizer_pt = torch.optim.AdamW(params=parameters_pt, lr=self.lr, weight_decay=cfg_opt.weight_decay_pt)
-            self.lr_scheduler_pt = get_scheduler(optimizer=self.optimizer_pt, **self.cfgs.train.scheduler_pt)
+            if isinstance(self.cfgs.train.scheduler_pt, partial):
+                self.lr_scheduler_pt = self.cfgs.train.scheduler_pt(optimizer=self.optimizer_pt)
+            else:
+                self.lr_scheduler_pt = get_scheduler(optimizer=self.optimizer_pt, **self.cfgs.train.scheduler_pt)
 
     def train(self):
         total_batch_size = sum(self.batch_size_list)*self.world_size*self.cfgs.train.gradient_accumulation_steps
 
         self.loggers.info("***** Running training *****")
         self.loggers.info(f"  Num batches each epoch = {len(self.train_loader_group.loader_list[0])}")
-        self.loggers.info(f"  Num Steps = {self.cfgs.train.scheduler.num_training_steps}")
+        self.loggers.info(f"  Num Steps = {self.cfgs.train.train_steps}")
         self.loggers.info(f"  Instantaneous batch size per device = {sum(self.batch_size_list)}")
         self.loggers.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
         self.loggers.info(f"  Gradient Accumulation steps = {self.cfgs.train.gradient_accumulation_steps}")
@@ -380,13 +397,13 @@ class Trainer:
                     lr_model = self.lr_scheduler.get_last_lr()[0] if hasattr(self, 'lr_scheduler') else 0.
                     lr_word = self.lr_scheduler_pt.get_last_lr()[0] if hasattr(self, 'lr_scheduler_pt') else 0.
                     self.loggers.log(datas={
-                        'Step':{'format':'[{}/{}]', 'data':[self.global_step, self.cfgs.train.scheduler.num_training_steps]},
+                        'Step':{'format':'[{}/{}]', 'data':[self.global_step, self.cfgs.train.train_steps]},
                         'LR_model':{'format':'{:.2e}', 'data':[lr_model]},
                         'LR_word':{'format':'{:.2e}', 'data':[lr_word]},
                         'Loss':{'format':'{:.5f}', 'data':[loss_sum.mean()]},
                     }, step=self.global_step)
 
-            if self.global_step>=self.cfgs.train.scheduler.num_training_steps:
+            if self.global_step>=self.cfgs.train.train_steps:
                 break
 
         self.wait_for_everyone()

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff/utils/net_utils.py

@@ -1,12 +1,11 @@
 import os
-from typing import Optional, Union, Tuple, Dict, Callable
+from typing import Optional, Union
 
 import torch
+from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION, Optimizer
 from torch import nn
 from torch.optim import lr_scheduler
-from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION, Optimizer
 from transformers import PretrainedConfig
-from collections import OrderedDict
 
 class TEUnetWrapper(nn.Module):
     def __init__(self, unet, TE):
@@ -33,7 +32,7 @@ def get_scheduler(
     optimizer: Optimizer,
     num_warmup_steps: Optional[int] = None,
     num_training_steps: Optional[int] = None,
-    scheduler_kwargs = {},
+    scheduler_kwargs={},
 ):
     """
     Unified API to get any scheduler from its name.
@@ -64,11 +63,11 @@ def get_scheduler(
     if num_warmup_steps is None:
         raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
 
-    #One Cycle for super convergence
-    if name=='one_cycle':
+    # One Cycle for super convergence
+    if name == 'one_cycle':
         scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=[x['lr'] for x in optimizer.state_dict()['param_groups']],
-                                steps_per_epoch=num_training_steps, epochs=1,
-                                pct_start=num_warmup_steps/num_training_steps, **scheduler_kwargs)
+                                            steps_per_epoch=num_training_steps, epochs=1,
+                                            pct_start=num_warmup_steps/num_training_steps, **scheduler_kwargs)
         return scheduler
 
     name = SchedulerType(name)
@@ -117,44 +116,46 @@ def remove_all_hooks(model: nn.Module) -> None:
         child._backward_hooks.clear()
 
 def remove_layers(model: nn.Module, layer_class):
-    named_modules = {k: v for k, v in model.named_modules()}
-    for k,v in named_modules.items():
+    named_modules = {k:v for k, v in model.named_modules()}
+    for k, v in named_modules.items():
         if isinstance(v, layer_class):
             parent, name = named_modules[k.rsplit('.', 1)]
             delattr(parent, name)
             del v
 
 def load_emb(path):
-    emb=torch.load(path, map_location='cpu')['string_to_param']['*']
+    emb = torch.load(path, map_location='cpu')['string_to_param']['*']
     emb.requires_grad_(False)
     return emb
 
-def save_emb(path, emb:torch.Tensor, replace=False):
+def save_emb(path, emb: torch.Tensor, replace=False):
     name = os.path.basename(path)
     if os.path.exists(path) and not replace:
         raise FileExistsError(f'embedding "{name}" already exist.')
-    name=name[:name.rfind('.')]
+    name = name[:name.rfind('.')]
     torch.save({'string_to_param':{'*':emb}, 'name':name}, path)
 
-
 def hook_compile(model):
     named_modules = {k:v for k, v in model.named_modules()}
 
     for name, block in named_modules.items():
         if len(block._forward_hooks)>0:
-            for hook in block._forward_hooks.values(): # 从前往后执行
+            for hook in block._forward_hooks.values():  # 从前往后执行
                 old_forward = block.forward
+
                 def new_forward(*args, **kwargs):
                     result = old_forward(*args, **kwargs)
                     hook_result = hook(block, args, result)
                     if hook_result is not None:
                         result = hook_result
                     return result
+
                 block.forward = new_forward
 
         if len(block._forward_pre_hooks)>0:
-            for hook in list(block._forward_pre_hooks.values())[::-1]: # 从前往后执行
+            for hook in list(block._forward_pre_hooks.values())[::-1]:  # 从前往后执行
                 old_forward = block.forward
+
                 def new_forward(*args, **kwargs):
                     result = hook(block, args)
                     if result is not None:
@@ -163,5 +164,6 @@ def hook_compile(model):
                     else:
                         result = args
                     return old_forward(*result, **kwargs)
+
                 block.forward = new_forward
-    remove_all_hooks(model)
+    remove_all_hooks(model)

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff/visualizer.py

@@ -124,12 +124,9 @@ class Visualizer:
             images = self.pipe(prompt_embeds=emb_p, negative_prompt_embeds=emb_n, **kwargs).images
         return images
 
-    @torch.no_grad()
-    def vis_to_dir(self, root, prompt, negative_prompt='', save_cfg=True, **kwargs):
+    def save_images(self, images, root, prompt, negative_prompt='', save_cfg=True):
         os.makedirs(root, exist_ok=True)
-        num_img_exist = max([int(x.split('-',1)[0]) for x in os.listdir(root) if x.rsplit('.', 1)[-1] in types_support])+1
-
-        images = self.vis_images(prompt, negative_prompt, **kwargs)
+        num_img_exist = max([int(x.split('-', 1)[0]) for x in os.listdir(root) if x.rsplit('.', 1)[-1] in types_support]) + 1
 
         for p, pn, img in zip(prompt, negative_prompt, images):
             img.save(os.path.join(root, f"{num_img_exist}-{to_validate_file(prompt[0])}.{self.cfgs.save.image_type}"), quality=self.cfgs.save.quality)
@@ -139,6 +136,10 @@ class Visualizer:
                     f.write(OmegaConf.to_yaml(self.cfgs_raw))
             num_img_exist += 1
 
+    def vis_to_dir(self, root, prompt, negative_prompt='', save_cfg=True, **kwargs):
+        images = self.vis_images(prompt, negative_prompt, **kwargs)
+        self.save_images(images, root, prompt, negative_prompt, save_cfg=save_cfg)
+
     def show_latent(self, prompt, negative_prompt='', **kwargs):
         emb_n, emb_p = self.te_hook.encode_prompt_to_emb(negative_prompt + prompt).chunk(2)
         emb_p = self.te_hook.mult_attn(emb_p, self.token_ex.parse_attn_mult(prompt))

{hcpdiff-0.3.5 → hcpdiff-0.3.7/hcpdiff.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hcpdiff
-Version: 0.3.5
+Version: 0.3.7
 Summary: A universal Stable-Diffusion toolbox
 Home-page: https://github.com/7eu7d7/HCP-Diffusion
 Author: Ziyi Dong
@@ -58,6 +58,8 @@ Compared to DreamArtist, DreamArtist++ is more stable with higher image quality
 * safetensors support
 * Controlnet (support train)
 * Min-SNR loss
+* Custom optimizer (Lion, DAdaptation, pytorch-optimizer, ...)
+* Custom lr scheduler
 
 ## Install
 

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/hcpdiff.egg-info/SOURCES.txt

@@ -17,6 +17,7 @@ cfgs/train/examples/CustomDiffusion.yaml
 cfgs/train/examples/DreamArtist++.yaml
 cfgs/train/examples/DreamArtist.yaml
 cfgs/train/examples/DreamBooth.yaml
+cfgs/train/examples/Lion_optimizer.yaml
 cfgs/train/examples/TextualInversion.yaml
 cfgs/train/examples/controlnet.yaml
 cfgs/train/examples/fine-tuning.yaml

{hcpdiff-0.3.5 → hcpdiff-0.3.7}/setup.py

@@ -23,7 +23,7 @@ def get_data_files(data_dir, prefix=''):
 setuptools.setup(
     name="hcpdiff",
     py_modules=["hcpdiff"],
-    version="0.3.5",
+    version="0.3.7",
     author="Ziyi Dong",
     author_email="dzy7eu7d7@gmail.com",
     description="A universal Stable-Diffusion toolbox",

hcpdiff-0.3.5/cfgs/train/examples/min_snr.yaml

@@ -1,58 +0,0 @@
-_base_: [cfgs/train/train_base.yaml, cfgs/train/tuning_base.yaml]
-
-unet:
-  -
-    lr: 1e-6
-    layers:
-      - '' # fine-tuning all layers in unet
-
-## fine-tuning text-encoder
-#text_encoder:
-#  - lr: 1e-6
-#    layers:
-#      - ''
-
-tokenizer_pt:
-  train: null
-
-train:
-  gradient_accumulation_steps: 1
-  save_step: 100
-
-  loss:
-    criterion: # min SNR loss
-      _target_: hcpdiff.loss.MinSNRLoss
-      gamma: 2.0
-
-  scheduler:
-    name: 'constant_with_warmup'
-    num_warmup_steps: 50
-    num_training_steps: 600
-
-model:
-  pretrained_model_name_or_path: 'runwayml/stable-diffusion-v1-5'
-  tokenizer_repeats: 1
-  ema_unet: 0
-  ema_text_encoder: 0
-
-data:
-  dataset1:
-    batch_size: 4
-    cache_latents: True
-
-    source:
-      data_source1:
-        img_root: 'imgs/'
-        prompt_template: 'prompt_tuning_template/object.txt'
-        caption_file: null # path to image captions (file_words)
-        tag_transforms:
-          transforms:
-            - _target_: hcpdiff.utils.caption_tools.TagShuffle
-            - _target_: hcpdiff.utils.caption_tools.TagDropout
-              p: 0.1
-            - _target_: hcpdiff.utils.caption_tools.TemplateFill
-              word_names: { }
-    bucket:
-      _target_: hcpdiff.data.bucket.RatioBucket.from_files # aspect ratio bucket
-      target_area: {_target_: "builtins.eval", _args_: ['512*512']}
-      num_bucket: 5