ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 19:21:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed issue with prompt token replace adding more than one replacement

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-26 18:52:23 -06:00
parent fd338e67bb
commit 6bd3851058
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
jobs/process/BaseSDTrainProcess.py
View File

@@ -312,6 +312,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
if self.trigger_word is not None and not is_reg: if self.trigger_word is not None and not is_reg:
prompt = self.sd.inject_trigger_into_prompt( prompt = self.sd.inject_trigger_into_prompt(
prompt, prompt,
trigger=self.trigger_word,
add_if_not_present=True, add_if_not_present=True,
) )
conditioned_prompts.append(prompt) conditioned_prompts.append(prompt)

3
toolkit/stable_diffusion_model.py
View File

@@ -698,6 +698,7 @@ class StableDiffusion:
return prompt return prompt
output_prompt = prompt output_prompt = prompt
default_replacements = ["[name]", "[trigger]"] default_replacements = ["[name]", "[trigger]"]
num_times_trigger_exists = prompt.count(trigger)
replace_with = trigger replace_with = trigger
if to_replace_list is None: if to_replace_list is None:
@@ -714,7 +715,7 @@ class StableDiffusion:
output_prompt = output_prompt.replace(to_replace, replace_with) output_prompt = output_prompt.replace(to_replace, replace_with)
# see how many times replace_with is in the prompt # see how many times replace_with is in the prompt
num_instances = prompt.count(replace_with) num_instances = output_prompt.count(replace_with)
if num_instances == 0 and add_if_not_present: if num_instances == 0 and add_if_not_present:
# add it to the beginning of the prompt # add it to the beginning of the prompt