Made encode_from_tokens_scheduled work when no hooks are set on patcher

This commit is contained in:
Jedrzej Kosinski 2024-11-19 17:17:01 -06:00
parent 9b2b1303b7
commit 59891b0f96

View File

@ -131,6 +131,15 @@ class CLIP:
def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict: dict[str]={}, show_pbar=True):
all_cond_pooled: list[tuple[torch.Tensor, dict[str]]] = []
all_hooks = self.patcher.forced_hooks
if all_hooks is None:
# if no hooks, do unscheduled encode_from_tokens and perform add_dict
return_pooled = "unprojected" if unprojected else True
pooled_dict = self.encode_from_tokens(tokens, return_pooled=return_pooled, return_dict=True)
cond = pooled_dict.pop("cond")
# add/update any keys with the provided add_dict
pooled_dict.update(add_dict)
all_cond_pooled.append([cond, pooled_dict])
else:
scheduled_keyframes = all_hooks.get_hooks_for_clip_schedule()
self.cond_stage_model.reset_clip_options()
@ -194,8 +203,7 @@ class CLIP:
if len(o) > 2:
for k in o[2]:
out[k] = o[2][k]
if self.apply_hooks_to_conds:
out["hooks"] = self.apply_hooks_to_conds
self.add_hooks_to_dict(out)
return out
if return_pooled: