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

Re-enable app shutdown actions (#6244)

Browse Source 
* closes #6242

* only override sigINT during slow model scanning

* fix ruff formatting

---------

Co-authored-by: Lincoln Stein <lstein@gmail.com>
This commit is contained in:
Lincoln Stein 2024-04-19 06:45:42 -04:00 committed by GitHub
parent a35386f24c
commit 2b9f06dc4c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 32 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
invokeai/app/services/model_install/model_install_default.py
View File

@ -3,7 +3,6 @@
import locale import locale
import os import os
import re import re
import signal
import threading import threading
import time import time
from hashlib import sha256 from hashlib import sha256
@ -43,6 +42,7 @@ from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMet
from invokeai.backend.model_manager.probe import ModelProbe from invokeai.backend.model_manager.probe import ModelProbe
from invokeai.backend.model_manager.search import ModelSearch from invokeai.backend.model_manager.search import ModelSearch
from invokeai.backend.util import InvokeAILogger from invokeai.backend.util import InvokeAILogger
from invokeai.backend.util.catch_sigint import catch_sigint
from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.devices import TorchDevice
from .model_install_base import ( from .model_install_base import (
@ -112,17 +112,6 @@ class ModelInstallService(ModelInstallServiceBase):
def start(self, invoker: Optional[Invoker] = None) -> None: def start(self, invoker: Optional[Invoker] = None) -> None:
"""Start the installer thread.""" """Start the installer thread."""
# Yes, this is weird. When the installer thread is running, the
# thread masks the ^C signal. When we receive a
# sigINT, we stop the thread, reset sigINT, and send a new
# sigINT to the parent process.
def sigint_handler(signum, frame):
self.stop()
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.raise_signal(signal.SIGINT)
signal.signal(signal.SIGINT, sigint_handler)
with self._lock: with self._lock:
if self._running: if self._running:
raise Exception("Attempt to start the installer service twice") raise Exception("Attempt to start the installer service twice")
@ -132,6 +121,7 @@ class ModelInstallService(ModelInstallServiceBase):
# In normal use, we do not want to scan the models directory - it should never have orphaned models. # In normal use, we do not want to scan the models directory - it should never have orphaned models.
# We should only do the scan when the flag is set (which should only be set when testing). # We should only do the scan when the flag is set (which should only be set when testing).
if self.app_config.scan_models_on_startup: if self.app_config.scan_models_on_startup:
with catch_sigint():
self._register_orphaned_models() self._register_orphaned_models()
# Check all models' paths and confirm they exist. A model could be missing if it was installed on a volume # Check all models' paths and confirm they exist. A model could be missing if it was installed on a volume

29
invokeai/backend/util/catch_sigint.py Normal file
View File

@ -0,0 +1,29 @@
"""
This module defines a context manager `catch_sigint()` which temporarily replaces
the sigINT handler defined by the ASGI in order to allow the user to ^C the application
and shut it down immediately. This was implemented in order to allow the user to interrupt
slow model hashing during startup.
Use like this:
from invokeai.backend.util.catch_sigint import catch_sigint
with catch_sigint():
run_some_hard_to_interrupt_process()
"""
import signal
from contextlib import contextmanager
from typing import Generator
def sigint_handler(signum, frame): # type: ignore
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.raise_signal(signal.SIGINT)
@contextmanager
def catch_sigint() -> Generator[None, None, None]:
original_handler = signal.getsignal(signal.SIGINT)
signal.signal(signal.SIGINT, sigint_handler)
yield
signal.signal(signal.SIGINT, original_handler)