Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into stalker-modular_controlnet

Browse Source
This commit is contained in:
Ryan Dick
2024-07-23 10:19:12 -04:00
parent 39e804d0f8 7b8e25f525
commit e2e47fd606
53 changed files with 1504 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
invokeai/backend/stable_diffusion/extensions_manager.py
View File

@ -63,9 +63,13 @@ class ExtensionsManager:
yield None
@contextmanager
def patch_unet(self, state_dict: Dict[str, torch.Tensor], unet: UNet2DConditionModel):
def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None):
if self._is_canceled and self._is_canceled():
raise CanceledException
# TODO: create logic in PR with extension which uses it
yield None
# TODO: create weight patch logic in PR with extension which uses it
with ExitStack() as exit_stack:
for ext in self._extensions:
exit_stack.enter_context(ext.patch_unet(unet, cached_weights))
yield None