import os import io import json import logging import asyncio import functools from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from PIL import Image from rembg import remove, new_session import redis.asyncio as aioredis import httpx from dotenv import load_dotenv from fastapi import FastAPI from .config import MODEL_NAME, PREFIX load_dotenv() logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) REDIS_URL = os.getenv("REDIS_URL") WEB_APP_URL = os.getenv("WEB_APP_URL") if not REDIS_URL: logger.error("REDIS_URL is not set. Worker will fail to start.") exit(1) if not WEB_APP_URL: logger.error("WEB_APP_URL is not set. Worker will fail to upload.") exit(1) WORKER_SECRET = os.getenv("WORKER_SECRET") r = aioredis.from_url( REDIS_URL, decode_responses=True, socket_timeout=15.0, socket_connect_timeout=10.0, socket_keepalive=True, health_check_interval=15, retry_on_timeout=True, ) session = new_session(MODEL_NAME) cpu_semaphore = asyncio.Semaphore(1) cpu_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="cpu_worker") network_semaphore = asyncio.Semaphore(5) async def upload_to_uploadthing( client: httpx.AsyncClient, image_bytes: bytes, filename: str, job_id: str = None, user_id: str = None, source_url: str = None, original_filename: str = None, ) -> str: url = f"{WEB_APP_URL}/api/worker/upload" files = {"file": (filename, image_bytes, "image/png")} data = {} if job_id: data["jobId"] = job_id if user_id: data["userId"] = user_id if source_url: data["sourceUrl"] = source_url if original_filename: data["originalName"] = original_filename headers = {"Authorization": f"Bearer {WORKER_SECRET}"} if WORKER_SECRET else {} response = await client.post(url, data=data, files=files, headers=headers, timeout=60.0) if response.status_code != 200: logger.error(f"Worker Upload API error: {response.status_code} - {response.text}") response.raise_for_status() data_res = response.json() blob_url = data_res.get("url") if not blob_url: raise ValueError("Worker Upload API did not return a URL") logger.info(f"Successfully uploaded to Uploadthing: {blob_url}") return blob_url async def process_job(job_data_str: str, client: httpx.AsyncClient): job_id = None try: job = json.loads(job_data_str) job_id = job["id"] source_url = job["url"] original_filename = job.get("filename", "image.png") user_id = job.get("userId") logger.info(f"Processing job {job_id} from {source_url}") response = await client.get(source_url, timeout=30.0) response.raise_for_status() await r.hset(f"{PREFIX}:job_status:{job_id}", mapping={"status": "processing"}) loop = asyncio.get_running_loop() async with cpu_semaphore: input_image = await loop.run_in_executor( cpu_executor, Image.open, io.BytesIO(response.content) ) output_image = await loop.run_in_executor( cpu_executor, functools.partial(remove, session=session), input_image ) img_byte_arr = io.BytesIO() await loop.run_in_executor( cpu_executor, output_image.save, img_byte_arr, "PNG" ) img_bytes = img_byte_arr.getvalue() name_parts = original_filename.rsplit(".", 1) if len(name_parts) == 2: base_name, ext = name_parts filename = f"{base_name}-nobg.{ext}" else: filename = f"{original_filename}-nobg.png" result_url = await upload_to_uploadthing( client, img_bytes, filename, job_id=job_id, user_id=user_id, source_url=source_url, original_filename=original_filename, ) await r.hset( f"{PREFIX}:job_status:{job_id}", mapping={ "status": "completed", "result_url": result_url, "sourceUrl": source_url, "filename": original_filename, }, ) await r.expire(f"{PREFIX}:job_status:{job_id}", 3600) logger.info(f"Job {job_id} completed successfully. URL: {result_url}") except Exception as e: logger.error(f"Error processing job: {str(e)}") if job_id: try: await r.hset(f"{PREFIX}:job_status:{job_id}", mapping={"status": "failed"}) await r.expire(f"{PREFIX}:job_status:{job_id}", 3600) except Exception as redis_err: logger.error(f"Failed to update failed status in Redis: {str(redis_err)}") async def worker_loop(client: httpx.AsyncClient): logger.info(f"Starting Redis worker loop... Listening on queue: {PREFIX}:job_queue") while True: try: await network_semaphore.acquire() try: result = await r.brpop(f"{PREFIX}:job_queue", timeout=5) if result: _, job_data_str = result async def run_and_release(): try: await process_job(job_data_str, client) finally: network_semaphore.release() asyncio.create_task(run_and_release()) else: network_semaphore.release() await asyncio.sleep(0.1) except Exception as inner_e: network_semaphore.release() raise inner_e except asyncio.CancelledError: logger.info("Worker loop cancelled.") break except (aioredis.TimeoutError, aioredis.ConnectionError, asyncio.TimeoutError) as te: logger.warning(f"Redis connection timeout/reset (reconnecting...): {str(te)}") await asyncio.sleep(2) except Exception as e: logger.error(f"Redis connection/worker loop error: {str(e)}") await asyncio.sleep(5) http_client: httpx.AsyncClient | None = None @asynccontextmanager async def lifespan(app: FastAPI): global http_client http_client = httpx.AsyncClient(timeout=60.0) worker_task = asyncio.create_task(worker_loop(http_client)) yield worker_task.cancel() try: await worker_task except asyncio.CancelledError: pass await http_client.aclose() cpu_executor.shutdown(wait=True) app = FastAPI(title="NoBG Worker", lifespan=lifespan) @app.get("/") def health_check(): return {"status": "Ok", "message": "NoBG Worker is Running"} if __name__ == "__main__": import uvicorn port = int(os.getenv("PORT", 8000)) uvicorn.run(app, host="0.0.0.0", port=port)