Spaces:
Sleeping
Sleeping
import gc | |
import threading | |
import time | |
import uuid | |
from typing import Optional | |
from loguru import logger | |
from carvekit.api.interface import Interface | |
from carvekit.web.schemas.config import WebAPIConfig | |
from carvekit.web.utils.init_utils import init_interface | |
from carvekit.web.other.removebg import process_remove_bg | |
class MLProcessor(threading.Thread): | |
"""Simple ml task queue processor""" | |
def __init__(self, api_config: WebAPIConfig): | |
super().__init__() | |
self.api_config = api_config | |
self.interface: Optional[Interface] = None | |
self.jobs = {} | |
self.completed_jobs = {} | |
def run(self): | |
"""Starts listening for new jobs.""" | |
unused_completed_jobs_timer = time.time() | |
if self.interface is None: | |
self.interface = init_interface(self.api_config) | |
while True: | |
# Clear unused completed jobs every hour | |
if time.time() - unused_completed_jobs_timer > 60: | |
self.clear_old_completed_jobs() | |
unused_completed_jobs_timer = time.time() | |
if len(self.jobs.keys()) >= 1: | |
id = list(self.jobs.keys())[0] | |
data = self.jobs[id] | |
# TODO add pydantic scheme here | |
response = process_remove_bg( | |
self.interface, data[0], data[1], data[2], data[3] | |
) | |
self.completed_jobs[id] = [response, time.time()] | |
try: | |
del self.jobs[id] | |
except KeyError or NameError as e: | |
logger.error(f"Something went wrong with Task Queue: {str(e)}") | |
gc.collect() | |
else: | |
time.sleep(1) | |
continue | |
def clear_old_completed_jobs(self): | |
"""Clears old completed jobs""" | |
if len(self.completed_jobs.keys()) >= 1: | |
for job_id in self.completed_jobs.keys(): | |
job_finished_time = self.completed_jobs[job_id][1] | |
if time.time() - job_finished_time > 3600: | |
try: | |
del self.completed_jobs[job_id] | |
except KeyError or NameError as e: | |
logger.error(f"Something went wrong with Task Queue: {str(e)}") | |
gc.collect() | |
def job_status(self, id: str) -> str: | |
""" | |
Returns current job status | |
Args: | |
id: id of the job | |
Returns: | |
Current job status for specified id. Job status can be [finished, wait, not_found] | |
""" | |
if id in self.completed_jobs.keys(): | |
return "finished" | |
elif id in self.jobs.keys(): | |
return "wait" | |
else: | |
return "not_found" | |
def job_result(self, id: str): | |
""" | |
Returns job processing result. | |
Args: | |
id: id of the job | |
Returns: | |
job processing result. | |
""" | |
if id in self.completed_jobs.keys(): | |
data = self.completed_jobs[id][0] | |
try: | |
del self.completed_jobs[id] | |
except KeyError or NameError: | |
pass | |
return data | |
else: | |
return False | |
def job_create(self, data: list): | |
""" | |
Send job to ML Processor | |
Args: | |
data: data object | |
""" | |
if self.is_alive() is False: | |
self.start() | |
id = uuid.uuid4().hex | |
self.jobs[id] = data | |
return id | |