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"},
|
"sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
|
||||||
},
|
},
|
||||||
opset=ctx.opset,
|
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
|
# VAE DECODER
|
||||||
|
|
|
@ -1,5 +1,6 @@
|
||||||
import atexit
|
import atexit
|
||||||
import gc
|
import gc
|
||||||
|
from functools import partial
|
||||||
from logging import getLogger
|
from logging import getLogger
|
||||||
|
|
||||||
from diffusers.utils.logging import disable_progress_bar
|
from diffusers.utils.logging import disable_progress_bar
|
||||||
|
@ -67,11 +68,11 @@ def run():
|
||||||
app, pool = main()
|
app, pool = main()
|
||||||
pool.start()
|
pool.start()
|
||||||
|
|
||||||
def quit():
|
def quit(p: DevicePoolExecutor):
|
||||||
logger.info("shutting down workers")
|
logger.info("shutting down workers")
|
||||||
pool.join()
|
p.join()
|
||||||
|
|
||||||
atexit.register(quit)
|
atexit.register(partial(quit, pool))
|
||||||
return app
|
return app
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue