1 files changed, 6 insertions, 6 deletions
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 720ebf3..289d6bd 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -72,7 +72,7 @@ def textual_inversion_strategy_callbacks(
    if use_ema:
        ema_embeddings = EMAModel(
-            text_encoder.text_model.embeddings.token_override_embedding.params.parameters(),
+            text_encoder.text_model.embeddings.token_override_embedding.parameters(),
            inv_gamma=ema_inv_gamma,
            power=ema_power,
            max_value=ema_max_decay,
@@ -84,20 +84,20 @@ def textual_inversion_strategy_callbacks(
    def ema_context():
        if ema_embeddings is not None:
            return ema_embeddings.apply_temporary(
-                text_encoder.text_model.embeddings.token_override_embedding.params.parameters()
+                text_encoder.text_model.embeddings.token_override_embedding.parameters()
            )
        else:
            return nullcontext()
    @contextmanager
    def on_train(epoch: int):
-        text_encoder.text_model.embeddings.token_override_embedding.params.train()
+        text_encoder.train()
        tokenizer.train()
        yield
    @contextmanager
    def on_eval():
-        text_encoder.text_model.embeddings.token_override_embedding.params.eval()
+        text_encoder.eval()
        tokenizer.eval()
        with ema_context():
@@ -108,7 +108,7 @@ def textual_inversion_strategy_callbacks(
        if use_emb_decay:
            params = [
                p
-                for p in text_encoder.text_model.embeddings.token_override_embedding.params
+                for p in text_encoder.text_model.embeddings.token_override_embedding.parameters()
                if p.grad is not None
            ]
            return torch.stack(params) if len(params) != 0 else None
@@ -116,7 +116,7 @@ def textual_inversion_strategy_callbacks(
    @torch.no_grad()
    def on_after_optimize(w, lrs: dict[str, float]):
        if ema_embeddings is not None:
-            ema_embeddings.step(text_encoder.text_model.embeddings.token_override_embedding.params.parameters())
+            ema_embeddings.step(text_encoder.text_model.embeddings.token_override_embedding.parameters())
        if use_emb_decay and w is not None:
            lr = lrs["emb"] or lrs["0"]

diff --git a/training/strategy/ti.py b/training/strategy/ti.py index 720ebf3..289d6bd 100644 --- a/training/strategy/ti.py +++ b/training/strategy/ti.py
@@ -72,7 +72,7 @@ def textual_inversion_strategy_callbacks(
72		72
73	if use_ema:	73	if use_ema:
74	ema_embeddings = EMAModel(	74	ema_embeddings = EMAModel(
75	text_encoder.text_model.embeddings.token_override_embedding.params.parameters(),	75	text_encoder.text_model.embeddings.token_override_embedding.parameters(),
76	inv_gamma=ema_inv_gamma,	76	inv_gamma=ema_inv_gamma,
77	power=ema_power,	77	power=ema_power,
78	max_value=ema_max_decay,	78	max_value=ema_max_decay,
@@ -84,20 +84,20 @@ def textual_inversion_strategy_callbacks(
84	def ema_context():	84	def ema_context():
85	if ema_embeddings is not None:	85	if ema_embeddings is not None:
86	return ema_embeddings.apply_temporary(	86	return ema_embeddings.apply_temporary(
87	text_encoder.text_model.embeddings.token_override_embedding.params.parameters()	87	text_encoder.text_model.embeddings.token_override_embedding.parameters()
88	)	88	)
89	else:	89	else:
90	return nullcontext()	90	return nullcontext()
91		91
92	@contextmanager	92	@contextmanager
93	def on_train(epoch: int):	93	def on_train(epoch: int):
94	text_encoder.text_model.embeddings.token_override_embedding.params.train()	94	text_encoder.train()
95	tokenizer.train()	95	tokenizer.train()
96	yield	96	yield
97		97
98	@contextmanager	98	@contextmanager
99	def on_eval():	99	def on_eval():
100	text_encoder.text_model.embeddings.token_override_embedding.params.eval()	100	text_encoder.eval()
101	tokenizer.eval()	101	tokenizer.eval()
102		102
103	with ema_context():	103	with ema_context():
@@ -108,7 +108,7 @@ def textual_inversion_strategy_callbacks(
108	if use_emb_decay:	108	if use_emb_decay:
109	params = [	109	params = [
110	p	110	p
111	for p in text_encoder.text_model.embeddings.token_override_embedding.params	111	for p in text_encoder.text_model.embeddings.token_override_embedding.parameters()
112	if p.grad is not None	112	if p.grad is not None
113	]	113	]
114	return torch.stack(params) if len(params) != 0 else None	114	return torch.stack(params) if len(params) != 0 else None
@@ -116,7 +116,7 @@ def textual_inversion_strategy_callbacks(
116	@torch.no_grad()	116	@torch.no_grad()
117	def on_after_optimize(w, lrs: dict[str, float]):	117	def on_after_optimize(w, lrs: dict[str, float]):
118	if ema_embeddings is not None:	118	if ema_embeddings is not None:
119	ema_embeddings.step(text_encoder.text_model.embeddings.token_override_embedding.params.parameters())	119	ema_embeddings.step(text_encoder.text_model.embeddings.token_override_embedding.parameters())
120		120
121	if use_emb_decay and w is not None:	121	if use_emb_decay and w is not None:
122	lr = lrs["emb"] or lrs["0"]	122	lr = lrs["emb"] or lrs["0"]