| 36 | self._workers.append(t) |
| 37 | |
| 38 | def _worker_loop(self, device: torch.device, work_q: queue.Queue, result_q: queue.Queue): |
| 39 | try: |
| 40 | comfy.model_management.set_torch_device(device) |
| 41 | except Exception as e: |
| 42 | logging.error(f"MultiGPUThreadPool: failed to set device {device}: {e}") |
| 43 | while True: |
| 44 | item = work_q.get() |
| 45 | if item is None: |
| 46 | return |
| 47 | result_q.put((None, e)) |
| 48 | return |
| 49 | while True: |
| 50 | item = work_q.get() |
| 51 | if item is None: |
| 52 | break |
| 53 | fn, args, kwargs = item |
| 54 | try: |
| 55 | result = fn(*args, **kwargs) |
| 56 | result_q.put((result, None)) |
| 57 | except comfy.model_management.InterruptProcessingException as e: |
| 58 | result_q.put((None, e)) |
| 59 | except Exception as e: |
| 60 | result_q.put((None, e)) |
| 61 | |
| 62 | def submit(self, device: torch.device, fn, *args, **kwargs): |
| 63 | self._work_queues[device].put((fn, args, kwargs)) |