diff options
author | Volpeon <git@volpeon.ink> | 2023-01-15 10:12:04 +0100 |
---|---|---|
committer | Volpeon <git@volpeon.ink> | 2023-01-15 10:12:04 +0100 |
commit | 34648b763fa60e3161a5b5f1243ed1b5c3b0188e (patch) | |
tree | 4c2b8104a8d1af26955561959591249d9281a02f | |
parent | Added functional trainer (diff) | |
download | textual-inversion-diff-34648b763fa60e3161a5b5f1243ed1b5c3b0188e.tar.gz textual-inversion-diff-34648b763fa60e3161a5b5f1243ed1b5c3b0188e.tar.bz2 textual-inversion-diff-34648b763fa60e3161a5b5f1243ed1b5c3b0188e.zip |
Added functional TI strategy
-rw-r--r-- | train_ti.py | 108 | ||||
-rw-r--r-- | training/functional.py | 118 | ||||
-rw-r--r-- | training/strategy/ti.py | 164 |
3 files changed, 312 insertions, 78 deletions
diff --git a/train_ti.py b/train_ti.py index 97e4e72..2fd325b 100644 --- a/train_ti.py +++ b/train_ti.py | |||
@@ -17,7 +17,8 @@ from slugify import slugify | |||
17 | from util import load_config, load_embeddings_from_dir | 17 | from util import load_config, load_embeddings_from_dir |
18 | from data.csv import VlpnDataModule, VlpnDataItem | 18 | from data.csv import VlpnDataModule, VlpnDataItem |
19 | from trainer_old.base import Checkpointer | 19 | from trainer_old.base import Checkpointer |
20 | from training.functional import train, loss_step, train_loop, generate_class_images, add_placeholder_tokens, get_models | 20 | from training.functional import train, generate_class_images, add_placeholder_tokens, get_models |
21 | from training.strategy.ti import textual_inversion_strategy | ||
21 | from training.optimization import get_scheduler | 22 | from training.optimization import get_scheduler |
22 | from training.lr import LRFinder | 23 | from training.lr import LRFinder |
23 | from training.util import EMAModel, save_args | 24 | from training.util import EMAModel, save_args |
@@ -387,6 +388,11 @@ def parse_args(): | |||
387 | help="The weight of prior preservation loss." | 388 | help="The weight of prior preservation loss." |
388 | ) | 389 | ) |
389 | parser.add_argument( | 390 | parser.add_argument( |
391 | "--use_emb_decay", | ||
392 | action="store_true", | ||
393 | help="Whether to use embedding decay." | ||
394 | ) | ||
395 | parser.add_argument( | ||
390 | "--emb_decay_target", | 396 | "--emb_decay_target", |
391 | default=0.4, | 397 | default=0.4, |
392 | type=float, | 398 | type=float, |
@@ -591,14 +597,6 @@ def main(): | |||
591 | else: | 597 | else: |
592 | ema_embeddings = None | 598 | ema_embeddings = None |
593 | 599 | ||
594 | vae.requires_grad_(False) | ||
595 | unet.requires_grad_(False) | ||
596 | |||
597 | text_encoder.text_model.encoder.requires_grad_(False) | ||
598 | text_encoder.text_model.final_layer_norm.requires_grad_(False) | ||
599 | text_encoder.text_model.embeddings.position_embedding.requires_grad_(False) | ||
600 | text_encoder.text_model.embeddings.token_embedding.requires_grad_(False) | ||
601 | |||
602 | if args.scale_lr: | 600 | if args.scale_lr: |
603 | args.learning_rate = ( | 601 | args.learning_rate = ( |
604 | args.learning_rate * args.gradient_accumulation_steps * | 602 | args.learning_rate * args.gradient_accumulation_steps * |
@@ -719,73 +717,36 @@ def main(): | |||
719 | seed=args.seed, | 717 | seed=args.seed, |
720 | ) | 718 | ) |
721 | 719 | ||
722 | def on_prepare(): | 720 | strategy = textual_inversion_strategy( |
723 | text_encoder.text_model.embeddings.temp_token_embedding.requires_grad_(True) | ||
724 | |||
725 | if args.gradient_checkpointing: | ||
726 | unet.train() | ||
727 | |||
728 | @contextmanager | ||
729 | def on_train(epoch: int): | ||
730 | try: | ||
731 | tokenizer.train() | ||
732 | yield | ||
733 | finally: | ||
734 | pass | ||
735 | |||
736 | @contextmanager | ||
737 | def on_eval(): | ||
738 | try: | ||
739 | tokenizer.eval() | ||
740 | |||
741 | ema_context = ema_embeddings.apply_temporary( | ||
742 | text_encoder.text_model.embeddings.temp_token_embedding.parameters()) if args.use_ema else nullcontext() | ||
743 | |||
744 | with ema_context: | ||
745 | yield | ||
746 | finally: | ||
747 | pass | ||
748 | |||
749 | @torch.no_grad() | ||
750 | def on_after_optimize(lr: float): | ||
751 | if args.emb_decay_factor != 0: | ||
752 | text_encoder.text_model.embeddings.normalize( | ||
753 | args.emb_decay_target, | ||
754 | min(1.0, max(0.0, args.emb_decay_factor * ((lr - args.emb_decay_start) / (args.learning_rate - args.emb_decay_start)))) | ||
755 | ) | ||
756 | |||
757 | if args.use_ema: | ||
758 | ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters()) | ||
759 | |||
760 | def on_log(): | ||
761 | if args.use_ema: | ||
762 | return {"ema_decay": ema_embeddings.decay} | ||
763 | return {} | ||
764 | |||
765 | checkpointer = TextualInversionCheckpointer( | ||
766 | dtype=weight_dtype, | ||
767 | train_dataloader=train_dataloader, | ||
768 | val_dataloader=val_dataloader, | ||
769 | accelerator=accelerator, | 721 | accelerator=accelerator, |
770 | vae=vae, | ||
771 | unet=unet, | 722 | unet=unet, |
772 | tokenizer=tokenizer, | ||
773 | text_encoder=text_encoder, | 723 | text_encoder=text_encoder, |
774 | ema_embeddings=ema_embeddings, | 724 | tokenizer=tokenizer, |
725 | vae=vae, | ||
775 | sample_scheduler=sample_scheduler, | 726 | sample_scheduler=sample_scheduler, |
727 | train_dataloader=train_dataloader, | ||
728 | val_dataloader=val_dataloader, | ||
729 | dtype=weight_dtype, | ||
730 | output_dir=output_dir, | ||
731 | seed=args.seed, | ||
776 | placeholder_tokens=args.placeholder_tokens, | 732 | placeholder_tokens=args.placeholder_tokens, |
777 | placeholder_token_ids=placeholder_token_ids, | 733 | placeholder_token_ids=placeholder_token_ids, |
778 | output_dir=output_dir, | 734 | learning_rate=args.learning_rate, |
779 | sample_steps=args.sample_steps, | 735 | gradient_checkpointing=args.gradient_checkpointing, |
780 | sample_image_size=args.sample_image_size, | 736 | use_emb_decay=args.use_emb_decay, |
737 | emb_decay_target=args.emb_decay_target, | ||
738 | emb_decay_factor=args.emb_decay_factor, | ||
739 | emb_decay_start=args.emb_decay_start, | ||
740 | use_ema=args.use_ema, | ||
741 | ema_inv_gamma=args.ema_inv_gamma, | ||
742 | ema_power=args.ema_power, | ||
743 | ema_max_decay=args.ema_max_decay, | ||
781 | sample_batch_size=args.sample_batch_size, | 744 | sample_batch_size=args.sample_batch_size, |
782 | sample_batches=args.sample_batches, | 745 | sample_num_batches=args.sample_batches, |
783 | seed=args.seed | 746 | sample_num_steps=args.sample_steps, |
747 | sample_image_size=args.sample_image_size, | ||
784 | ) | 748 | ) |
785 | 749 | ||
786 | if accelerator.is_main_process: | ||
787 | accelerator.init_trackers("textual_inversion") | ||
788 | |||
789 | if args.find_lr: | 750 | if args.find_lr: |
790 | lr_finder = LRFinder( | 751 | lr_finder = LRFinder( |
791 | accelerator=accelerator, | 752 | accelerator=accelerator, |
@@ -793,10 +754,7 @@ def main(): | |||
793 | model=text_encoder, | 754 | model=text_encoder, |
794 | train_dataloader=train_dataloader, | 755 | train_dataloader=train_dataloader, |
795 | val_dataloader=val_dataloader, | 756 | val_dataloader=val_dataloader, |
796 | loss_step=loss_step_, | 757 | **strategy, |
797 | on_train=on_train, | ||
798 | on_eval=on_eval, | ||
799 | on_after_optimize=on_after_optimize, | ||
800 | ) | 758 | ) |
801 | lr_finder.run(num_epochs=100, end_lr=1e3) | 759 | lr_finder.run(num_epochs=100, end_lr=1e3) |
802 | 760 | ||
@@ -811,13 +769,7 @@ def main(): | |||
811 | checkpoint_frequency=args.checkpoint_frequency, | 769 | checkpoint_frequency=args.checkpoint_frequency, |
812 | global_step_offset=global_step_offset, | 770 | global_step_offset=global_step_offset, |
813 | prior_loss_weight=args.prior_loss_weight, | 771 | prior_loss_weight=args.prior_loss_weight, |
814 | on_prepare=on_prepare, | 772 | **strategy, |
815 | on_log=on_log, | ||
816 | on_train=on_train, | ||
817 | on_after_optimize=on_after_optimize, | ||
818 | on_eval=on_eval, | ||
819 | on_sample=checkpointer.save_samples, | ||
820 | on_checkpoint=checkpointer.checkpoint, | ||
821 | ) | 773 | ) |
822 | 774 | ||
823 | 775 | ||
diff --git a/training/functional.py b/training/functional.py index 1f2ca6d..e54c9c8 100644 --- a/training/functional.py +++ b/training/functional.py | |||
@@ -2,6 +2,8 @@ import math | |||
2 | from contextlib import _GeneratorContextManager, nullcontext | 2 | from contextlib import _GeneratorContextManager, nullcontext |
3 | from typing import Callable, Any, Tuple, Union, Optional | 3 | from typing import Callable, Any, Tuple, Union, Optional |
4 | from functools import partial | 4 | from functools import partial |
5 | from pathlib import Path | ||
6 | import itertools | ||
5 | 7 | ||
6 | import torch | 8 | import torch |
7 | import torch.nn.functional as F | 9 | import torch.nn.functional as F |
@@ -26,6 +28,14 @@ def const(result=None): | |||
26 | return fn | 28 | return fn |
27 | 29 | ||
28 | 30 | ||
31 | def make_grid(images, rows, cols): | ||
32 | w, h = images[0].size | ||
33 | grid = Image.new('RGB', size=(cols*w, rows*h)) | ||
34 | for i, image in enumerate(images): | ||
35 | grid.paste(image, box=(i % cols*w, i//cols*h)) | ||
36 | return grid | ||
37 | |||
38 | |||
29 | def get_models(pretrained_model_name_or_path: str): | 39 | def get_models(pretrained_model_name_or_path: str): |
30 | tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer') | 40 | tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer') |
31 | text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder') | 41 | text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder') |
@@ -40,6 +50,107 @@ def get_models(pretrained_model_name_or_path: str): | |||
40 | return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings | 50 | return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings |
41 | 51 | ||
42 | 52 | ||
53 | def save_samples( | ||
54 | accelerator: Accelerator, | ||
55 | unet: UNet2DConditionModel, | ||
56 | text_encoder: CLIPTextModel, | ||
57 | tokenizer: MultiCLIPTokenizer, | ||
58 | vae: AutoencoderKL, | ||
59 | sample_scheduler: DPMSolverMultistepScheduler, | ||
60 | train_dataloader: DataLoader, | ||
61 | val_dataloader: DataLoader, | ||
62 | dtype: torch.dtype, | ||
63 | output_dir: Path, | ||
64 | seed: int, | ||
65 | step: int, | ||
66 | batch_size: int = 1, | ||
67 | num_batches: int = 1, | ||
68 | num_steps: int = 20, | ||
69 | guidance_scale: float = 7.5, | ||
70 | image_size: Optional[int] = None, | ||
71 | ): | ||
72 | print(f"Saving samples for step {step}...") | ||
73 | |||
74 | samples_path = output_dir.joinpath("samples") | ||
75 | |||
76 | grid_cols = min(batch_size, 4) | ||
77 | grid_rows = (num_batches * batch_size) // grid_cols | ||
78 | |||
79 | unet = accelerator.unwrap_model(unet) | ||
80 | text_encoder = accelerator.unwrap_model(text_encoder) | ||
81 | |||
82 | orig_unet_dtype = unet.dtype | ||
83 | orig_text_encoder_dtype = text_encoder.dtype | ||
84 | |||
85 | unet.to(dtype=dtype) | ||
86 | text_encoder.to(dtype=dtype) | ||
87 | |||
88 | pipeline = VlpnStableDiffusion( | ||
89 | text_encoder=text_encoder, | ||
90 | vae=vae, | ||
91 | unet=unet, | ||
92 | tokenizer=tokenizer, | ||
93 | scheduler=sample_scheduler, | ||
94 | ).to(accelerator.device) | ||
95 | pipeline.set_progress_bar_config(dynamic_ncols=True) | ||
96 | |||
97 | generator = torch.Generator(device=accelerator.device).manual_seed(seed) | ||
98 | |||
99 | for pool, data, gen in [ | ||
100 | ("stable", val_dataloader, generator), | ||
101 | ("val", val_dataloader, None), | ||
102 | ("train", train_dataloader, None) | ||
103 | ]: | ||
104 | all_samples = [] | ||
105 | file_path = samples_path.joinpath(pool, f"step_{step}.jpg") | ||
106 | file_path.parent.mkdir(parents=True, exist_ok=True) | ||
107 | |||
108 | batches = list(itertools.islice(itertools.cycle(data), batch_size * num_batches)) | ||
109 | prompt_ids = [ | ||
110 | prompt | ||
111 | for batch in batches | ||
112 | for prompt in batch["prompt_ids"] | ||
113 | ] | ||
114 | nprompt_ids = [ | ||
115 | prompt | ||
116 | for batch in batches | ||
117 | for prompt in batch["nprompt_ids"] | ||
118 | ] | ||
119 | |||
120 | for i in range(num_batches): | ||
121 | start = i * batch_size | ||
122 | end = (i + 1) * batch_size | ||
123 | prompt = prompt_ids[start:end] | ||
124 | nprompt = nprompt_ids[start:end] | ||
125 | |||
126 | samples = pipeline( | ||
127 | prompt=prompt, | ||
128 | negative_prompt=nprompt, | ||
129 | height=image_size, | ||
130 | width=image_size, | ||
131 | generator=gen, | ||
132 | guidance_scale=guidance_scale, | ||
133 | num_inference_steps=num_steps, | ||
134 | output_type='pil' | ||
135 | ).images | ||
136 | |||
137 | all_samples += samples | ||
138 | |||
139 | image_grid = make_grid(all_samples, grid_rows, grid_cols) | ||
140 | image_grid.save(file_path, quality=85) | ||
141 | |||
142 | unet.to(dtype=orig_unet_dtype) | ||
143 | text_encoder.to(dtype=orig_text_encoder_dtype) | ||
144 | |||
145 | del unet | ||
146 | del text_encoder | ||
147 | del generator | ||
148 | del pipeline | ||
149 | |||
150 | if torch.cuda.is_available(): | ||
151 | torch.cuda.empty_cache() | ||
152 | |||
153 | |||
43 | def generate_class_images( | 154 | def generate_class_images( |
44 | accelerator: Accelerator, | 155 | accelerator: Accelerator, |
45 | text_encoder: CLIPTextModel, | 156 | text_encoder: CLIPTextModel, |
@@ -109,6 +220,10 @@ def get_models(pretrained_model_name_or_path: str): | |||
109 | 220 | ||
110 | embeddings = patch_managed_embeddings(text_encoder) | 221 | embeddings = patch_managed_embeddings(text_encoder) |
111 | 222 | ||
223 | vae.requires_grad_(False) | ||
224 | unet.requires_grad_(False) | ||
225 | text_encoder.requires_grad_(False) | ||
226 | |||
112 | return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings | 227 | return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings |
113 | 228 | ||
114 | 229 | ||
@@ -427,6 +542,9 @@ def train( | |||
427 | seed, | 542 | seed, |
428 | ) | 543 | ) |
429 | 544 | ||
545 | if accelerator.is_main_process: | ||
546 | accelerator.init_trackers("textual_inversion") | ||
547 | |||
430 | train_loop( | 548 | train_loop( |
431 | accelerator=accelerator, | 549 | accelerator=accelerator, |
432 | optimizer=optimizer, | 550 | optimizer=optimizer, |
diff --git a/training/strategy/ti.py b/training/strategy/ti.py new file mode 100644 index 0000000..83dc566 --- /dev/null +++ b/training/strategy/ti.py | |||
@@ -0,0 +1,164 @@ | |||
1 | from contextlib import nullcontext | ||
2 | from typing import Optional | ||
3 | from functools import partial | ||
4 | from contextlib import contextmanager, nullcontext | ||
5 | from pathlib import Path | ||
6 | |||
7 | import torch | ||
8 | from torch.utils.data import DataLoader | ||
9 | |||
10 | from accelerate import Accelerator | ||
11 | from transformers import CLIPTextModel | ||
12 | from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepScheduler | ||
13 | |||
14 | from slugify import slugify | ||
15 | |||
16 | from models.clip.tokenizer import MultiCLIPTokenizer | ||
17 | from training.util import EMAModel | ||
18 | from training.functional import save_samples | ||
19 | |||
20 | |||
21 | def textual_inversion_strategy( | ||
22 | accelerator: Accelerator, | ||
23 | unet: UNet2DConditionModel, | ||
24 | text_encoder: CLIPTextModel, | ||
25 | tokenizer: MultiCLIPTokenizer, | ||
26 | vae: AutoencoderKL, | ||
27 | sample_scheduler: DPMSolverMultistepScheduler, | ||
28 | train_dataloader: DataLoader, | ||
29 | val_dataloader: DataLoader, | ||
30 | dtype: torch.dtype, | ||
31 | output_dir: Path, | ||
32 | seed: int, | ||
33 | placeholder_tokens: list[str], | ||
34 | placeholder_token_ids: list[list[int]], | ||
35 | learning_rate: float, | ||
36 | gradient_checkpointing: bool = False, | ||
37 | use_emb_decay: bool = False, | ||
38 | emb_decay_target: float = 0.4, | ||
39 | emb_decay_factor: float = 1, | ||
40 | emb_decay_start: float = 1e-4, | ||
41 | use_ema: bool = False, | ||
42 | ema_inv_gamma: float = 1.0, | ||
43 | ema_power: int = 1, | ||
44 | ema_max_decay: float = 0.9999, | ||
45 | sample_batch_size: int = 1, | ||
46 | sample_num_batches: int = 1, | ||
47 | sample_num_steps: int = 20, | ||
48 | sample_guidance_scale: float = 7.5, | ||
49 | sample_image_size: Optional[int] = None, | ||
50 | ): | ||
51 | save_samples_ = partial( | ||
52 | save_samples, | ||
53 | accelerator=accelerator, | ||
54 | unet=unet, | ||
55 | text_encoder=text_encoder, | ||
56 | tokenizer=tokenizer, | ||
57 | vae=vae, | ||
58 | sample_scheduler=sample_scheduler, | ||
59 | train_dataloader=train_dataloader, | ||
60 | val_dataloader=val_dataloader, | ||
61 | dtype=dtype, | ||
62 | output_dir=output_dir, | ||
63 | seed=seed, | ||
64 | batch_size=sample_batch_size, | ||
65 | num_batches=sample_num_batches, | ||
66 | num_steps=sample_num_steps, | ||
67 | guidance_scale=sample_guidance_scale, | ||
68 | image_size=sample_image_size, | ||
69 | ) | ||
70 | |||
71 | if use_ema: | ||
72 | ema_embeddings = EMAModel( | ||
73 | text_encoder.text_model.embeddings.temp_token_embedding.parameters(), | ||
74 | inv_gamma=ema_inv_gamma, | ||
75 | power=ema_power, | ||
76 | max_value=ema_max_decay, | ||
77 | ) | ||
78 | else: | ||
79 | ema_embeddings = None | ||
80 | |||
81 | def on_prepare(): | ||
82 | text_encoder.text_model.embeddings.temp_token_embedding.requires_grad_(True) | ||
83 | |||
84 | if use_ema: | ||
85 | ema_embeddings.to(accelerator.device) | ||
86 | |||
87 | if gradient_checkpointing: | ||
88 | unet.train() | ||
89 | |||
90 | @contextmanager | ||
91 | def on_train(epoch: int): | ||
92 | try: | ||
93 | tokenizer.train() | ||
94 | yield | ||
95 | finally: | ||
96 | pass | ||
97 | |||
98 | @contextmanager | ||
99 | def on_eval(): | ||
100 | try: | ||
101 | tokenizer.eval() | ||
102 | |||
103 | ema_context = ema_embeddings.apply_temporary( | ||
104 | text_encoder.text_model.embeddings.temp_token_embedding.parameters()) if use_ema else nullcontext() | ||
105 | |||
106 | with ema_context: | ||
107 | yield | ||
108 | finally: | ||
109 | pass | ||
110 | |||
111 | @torch.no_grad() | ||
112 | def on_after_optimize(lr: float): | ||
113 | if use_emb_decay: | ||
114 | text_encoder.text_model.embeddings.normalize( | ||
115 | emb_decay_target, | ||
116 | min(1.0, max(0.0, emb_decay_factor * ((lr - emb_decay_start) / (learning_rate - emb_decay_start)))) | ||
117 | ) | ||
118 | |||
119 | if use_ema: | ||
120 | ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters()) | ||
121 | |||
122 | def on_log(): | ||
123 | if use_ema: | ||
124 | return {"ema_decay": ema_embeddings.decay} | ||
125 | return {} | ||
126 | |||
127 | @torch.no_grad() | ||
128 | def on_checkpoint(step, postfix): | ||
129 | print(f"Saving checkpoint for step {step}...") | ||
130 | |||
131 | checkpoints_path = output_dir.joinpath("checkpoints") | ||
132 | checkpoints_path.mkdir(parents=True, exist_ok=True) | ||
133 | |||
134 | text_encoder = accelerator.unwrap_model(text_encoder) | ||
135 | |||
136 | ema_context = ema_embeddings.apply_temporary( | ||
137 | text_encoder.text_model.embeddings.temp_token_embedding.parameters() | ||
138 | ) if ema_embeddings is not None else nullcontext() | ||
139 | |||
140 | with ema_context: | ||
141 | for (token, ids) in zip(placeholder_tokens, placeholder_token_ids): | ||
142 | text_encoder.text_model.embeddings.save_embed( | ||
143 | ids, | ||
144 | checkpoints_path.joinpath(f"{slugify(token)}_{step}_{postfix}.bin") | ||
145 | ) | ||
146 | |||
147 | @torch.no_grad() | ||
148 | def on_sample(step): | ||
149 | ema_context = ema_embeddings.apply_temporary( | ||
150 | text_encoder.text_model.embeddings.temp_token_embedding.parameters() | ||
151 | ) if ema_embeddings is not None else nullcontext() | ||
152 | |||
153 | with ema_context: | ||
154 | save_samples_(step=step) | ||
155 | |||
156 | return { | ||
157 | "on_prepare": on_prepare, | ||
158 | "on_train": on_train, | ||
159 | "on_eval": on_eval, | ||
160 | "on_after_optimize": on_after_optimize, | ||
161 | "on_log": on_log, | ||
162 | "on_checkpoint": on_checkpoint, | ||
163 | "on_sample": on_sample, | ||
164 | } | ||