RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: v5 fix 1 (#10) 

Co-authored-by: James <1561862923@qq.com>
This commit is contained in:
petmycat 2026-08-15 09:58:03 +08:00 committed by GitHub
parent aadec68812
commit ed9c9b9034
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 4 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
extensions_built_in/sd_trainer/SDTrainer.py
View File

@ -1930,9 +1930,8 @@ class SDTrainer(BaseSDTrainProcess):
encoded = {}
for source_name, prompts in source_prompts.items():
with self._activator_mode('full'):
active_embeds = self.sd.encode_prompt(
active_embeds = self.sd.get_prompt_embeds(
prompts,
long_prompts=self.do_long_prompts,
return_taps=context_config.enabled,
)
active_pred = self.predict_noise(
@ -1946,9 +1945,8 @@ class SDTrainer(BaseSDTrainProcess):
)
with self._activator_mode('activator_bypass'):
with torch.no_grad():
bypass_embeds = self.sd.encode_prompt(
bypass_embeds = self.sd.get_prompt_embeds(
prompts,
long_prompts=self.do_long_prompts,
return_taps=context_config.enabled,
).detach()
bypass_pred = self.predict_noise(

4
testing/test_sd_trainer_three_phase_runtime.py
View File

@ -331,7 +331,7 @@ class ThreePhaseRuntimeTest(unittest.TestCase):
trainer.do_long_prompts = False
trainer.additional_logs = {}
trainer.sd = SimpleNamespace(
encode_prompt=lambda prompts, **kwargs: _FakeEmbeds(len(prompts)),
get_prompt_embeds=lambda prompts, **kwargs: _FakeEmbeds(len(prompts)),
get_loss_target=lambda noise, batch, timesteps: noise + 1,
)
calls = []
@ -397,7 +397,7 @@ class ThreePhaseRuntimeTest(unittest.TestCase):
return embeds
trainer.sd = SimpleNamespace(
encode_prompt=encode,
get_prompt_embeds=encode,
get_loss_target=lambda noise, batch, timesteps: noise,
)
trainer.predict_noise = lambda noisy_latents, **kwargs: (