Fix batch & optimize

This commit is contained in:
AI-Casanova
2023-12-17 19:37:05 -06:00
parent 9e7757b7fe
commit 0b3e1f475d
3 changed files with 24 additions and 15 deletions
+14 -8
View File
@@ -58,33 +58,39 @@ class DiffusersTextualInversionManager(BaseTextualInversionManager):
debug(f'Prompt: expand={prompt}')
return self.pipe.tokenizer.encode(prompt, add_special_tokens=False)
def get_prompt_schedule(prompt, steps):
def get_prompt_schedule(p, prompt, steps):
temp = []
schedule = prompt_parser.get_learned_conditioning_prompt_schedules([prompt], steps)[0]
for chunk in schedule:
for s in range(steps):
if len(temp) < s + 1 <= chunk[0]:
temp.append(chunk[1])
return temp
return temp, len(schedule) > 1
def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, step: int = 1, clip_skip: typing.Optional[int] = None):
if 'StableDiffusion' not in pipe.__class__.__name__ and 'DemoFusion':
shared.log.warning(f"Prompt parser not supported: {pipe.__class__.__name__}")
return None, None, None, None
else:
positive_schedule = get_prompt_schedule(prompts[0], steps)
negative_schedule = get_prompt_schedule(negative_prompts[0], steps)
positive_schedule, scheduled = get_prompt_schedule(p, prompts[0], steps)
negative_schedule, neg_scheduled = get_prompt_schedule(p, negative_prompts[0], steps)
p.scheduled_prompt = scheduled or neg_scheduled
p.prompt_embeds = []
p.positive_pooleds = []
p.negative_embeds = []
p.negative_pooleds = []
cache = {}
for i in range(len(positive_schedule)):
prompt_embed, positive_pooled, negative_embed, negative_pooled = get_weighted_text_embeddings(pipe,
positive_schedule[i],
negative_schedule[i],
clip_skip)
cached = cache.get(positive_schedule[i]+negative_schedule[i], None)
if cached is not None:
prompt_embed, positive_pooled, negative_embed, negative_pooled = cached
else:
prompt_embed, positive_pooled, negative_embed, negative_pooled = get_weighted_text_embeddings(pipe,
positive_schedule[i],
negative_schedule[i],
clip_skip)
if prompt_embed is not None:
p.prompt_embeds.append(torch.cat([prompt_embed]*len(prompts), dim=0))
if negative_embed is not None: