summaryrefslogtreecommitdiffstats
path: root/training/strategy
diff options
context:
space:
mode:
authorVolpeon <git@volpeon.ink>2023-01-19 09:04:39 +0100
committerVolpeon <git@volpeon.ink>2023-01-19 09:04:39 +0100
commit2469501c3951a9ed86c820cddf7b32144a4a1c8d (patch)
tree9820efaa12fd31670616c1fd9da3e6bb06580aaf /training/strategy
parentUpdate (diff)
downloadtextual-inversion-diff-2469501c3951a9ed86c820cddf7b32144a4a1c8d.tar.gz
textual-inversion-diff-2469501c3951a9ed86c820cddf7b32144a4a1c8d.tar.bz2
textual-inversion-diff-2469501c3951a9ed86c820cddf7b32144a4a1c8d.zip
Move Accelerator preparation into strategy
Diffstat (limited to 'training/strategy')
-rw-r--r--training/strategy/dreambooth.py14
-rw-r--r--training/strategy/ti.py22
2 files changed, 34 insertions, 2 deletions
diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index f57e736..1277939 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -6,6 +6,7 @@ from pathlib import Path
6import itertools 6import itertools
7 7
8import torch 8import torch
9import torch.nn as nn
9from torch.utils.data import DataLoader 10from torch.utils.data import DataLoader
10 11
11from accelerate import Accelerator 12from accelerate import Accelerator
@@ -186,7 +187,18 @@ def dreambooth_strategy_callbacks(
186 ) 187 )
187 188
188 189
190def dreambooth_prepare(
191 accelerator: Accelerator,
192 text_encoder: CLIPTextModel,
193 unet: UNet2DConditionModel,
194 *args
195):
196 prep = [text_encoder, unet] + list(args)
197 text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(*prep)
198 return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler
199
200
189dreambooth_strategy = TrainingStrategy( 201dreambooth_strategy = TrainingStrategy(
190 callbacks=dreambooth_strategy_callbacks, 202 callbacks=dreambooth_strategy_callbacks,
191 prepare_unet=True 203 prepare=dreambooth_prepare
192) 204)
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index e922954..6a76f98 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -5,6 +5,7 @@ from contextlib import contextmanager, nullcontext
5from pathlib import Path 5from pathlib import Path
6 6
7import torch 7import torch
8import torch.nn as nn
8from torch.utils.data import DataLoader 9from torch.utils.data import DataLoader
9 10
10from accelerate import Accelerator 11from accelerate import Accelerator
@@ -94,7 +95,7 @@ def textual_inversion_strategy_callbacks(
94 return nullcontext() 95 return nullcontext()
95 96
96 def on_model(): 97 def on_model():
97 return text_encoder 98 return text_encoder.text_model.embeddings.temp_token_embedding
98 99
99 def on_prepare(): 100 def on_prepare():
100 text_encoder.text_model.embeddings.temp_token_embedding.requires_grad_(True) 101 text_encoder.text_model.embeddings.temp_token_embedding.requires_grad_(True)
@@ -163,6 +164,25 @@ def textual_inversion_strategy_callbacks(
163 ) 164 )
164 165
165 166
167def textual_inversion_prepare(
168 accelerator: Accelerator,
169 text_encoder: CLIPTextModel,
170 unet: UNet2DConditionModel,
171 *args
172):
173 weight_dtype = torch.float32
174 if accelerator.state.mixed_precision == "fp16":
175 weight_dtype = torch.float16
176 elif accelerator.state.mixed_precision == "bf16":
177 weight_dtype = torch.bfloat16
178
179 prep = [text_encoder] + list(args)
180 text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(*prep)
181 unet.to(accelerator.device, dtype=weight_dtype)
182 return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler
183
184
166textual_inversion_strategy = TrainingStrategy( 185textual_inversion_strategy = TrainingStrategy(
167 callbacks=textual_inversion_strategy_callbacks, 186 callbacks=textual_inversion_strategy_callbacks,
187 prepare=textual_inversion_prepare,
168) 188)