lint(api): explicitly bind the device pool to shutdown callback
This commit is contained in:
parent
c2f8fb1d31
commit
2bbc5d877f
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
||||
|
||||
|
|
Loading…
Reference in New Issue