Skip to content

Commit 39e10d8

Browse files
committed
Add invocation cancellation logic to patchers
1 parent 83a86ab commit 39e10d8

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

invokeai/backend/stable_diffusion/extensions_manager.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,6 @@ def _regenerate_ordered_callbacks(self):
4444
self._ordered_callbacks[callback_type] = sorted(callbacks, key=lambda x: x.metadata.order)
4545

4646
def run_callback(self, callback_type: ExtensionCallbackType, ctx: DenoiseContext):
47-
# TODO: add to patchers too?
48-
# and if so, should it be only in beginning of function or in for loop
4947
if self._is_canceled and self._is_canceled():
5048
raise CanceledException
5149

@@ -55,6 +53,9 @@ def run_callback(self, callback_type: ExtensionCallbackType, ctx: DenoiseContext
5553

5654
@contextmanager
5755
def patch_extensions(self, context: DenoiseContext):
56+
if self._is_canceled and self._is_canceled():
57+
raise CanceledException
58+
5859
with ExitStack() as exit_stack:
5960
for ext in self._extensions:
6061
exit_stack.enter_context(ext.patch_extension(context))
@@ -63,5 +64,8 @@ def patch_extensions(self, context: DenoiseContext):
6364

6465
@contextmanager
6566
def patch_unet(self, state_dict: Dict[str, torch.Tensor], unet: UNet2DConditionModel):
67+
if self._is_canceled and self._is_canceled():
68+
raise CanceledException
69+
6670
# TODO: create logic in PR with extension which uses it
6771
yield None

0 commit comments

Comments
 (0)