1
0
Fork 0

lint(api): explicitly bind the device pool to shutdown callback

This commit is contained in:
Sean Sube 2023-03-27 17:13:45 -05:00
parent c2f8fb1d31
commit 2bbc5d877f
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 5 additions and 4 deletions

View File

@ -267,7 +267,7 @@ def convert_diffusion_diffusers(
"sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
},
opset=ctx.opset,
half=False, # https://github.com/ssube/onnx-web/issues/290
half=False, # https://github.com/ssube/onnx-web/issues/290
)
# VAE DECODER

View File

@ -1,5 +1,6 @@
import atexit
import gc
from functools import partial
from logging import getLogger
from diffusers.utils.logging import disable_progress_bar
@ -67,11 +68,11 @@ def run():
app, pool = main()
pool.start()
def quit():
def quit(p: DevicePoolExecutor):
logger.info("shutting down workers")
pool.join()
p.join()
atexit.register(quit)
atexit.register(partial(quit, pool))
return app