Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
iperov
GitHub Repository: iperov/deepfacelab
Path: blob/master/core/joblib/MPClassFuncOnDemand.py
628 views
1
import multiprocessing
2
from core.interact import interact as io
3
4
class MPClassFuncOnDemand():
5
def __init__(self, class_handle, class_func_name, **class_kwargs):
6
self.class_handle = class_handle
7
self.class_func_name = class_func_name
8
self.class_kwargs = class_kwargs
9
10
self.class_func = None
11
12
self.s2c = multiprocessing.Queue()
13
self.c2s = multiprocessing.Queue()
14
self.lock = multiprocessing.Lock()
15
16
io.add_process_messages_callback(self.io_callback)
17
18
def io_callback(self):
19
while not self.c2s.empty():
20
func_args, func_kwargs = self.c2s.get()
21
if self.class_func is None:
22
self.class_func = getattr( self.class_handle(**self.class_kwargs), self.class_func_name)
23
self.s2c.put ( self.class_func (*func_args, **func_kwargs) )
24
25
def __call__(self, *args, **kwargs):
26
with self.lock:
27
self.c2s.put ( (args, kwargs) )
28
return self.s2c.get()
29
30
def __getstate__(self):
31
return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}
32
33
34