Sync from GitHub via hub-sync

.dockerignore

@@ -0,0 +1,17 @@
+__pycache__/
+*.py[cod]
+.env
+.venv/
+venv/
+env/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+
+uploads/
+outputs/
+dump.rdb
+*.rdb
+
+.git/
+.vscode/

Dockerfile

@@ -17,8 +17,14 @@ WORKDIR /app
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Step 2: Enable fast downloads
-ENV HF_HUB_ENABLE_HF_TRANSFER=1
+# Step 2: Runtime defaults for Hugging Face Spaces
+ENV HF_HUB_ENABLE_HF_TRANSFER=1 \
+    DEBUG=false \
+    RESET_JOBS_ON_STARTUP=true \
+    JOB_TTL_SECONDS=3600 \
+    JOB_CLEANUP_INTERVAL_SECONDS=300 \
+    CELERY_BROKER_URL=redis://localhost:6379/0 \
+    CELERY_RESULT_BACKEND=redis://localhost:6379/1
 
 # Step 3. Download models for caching
 RUN python -c "from huggingface_hub import snapshot_download; snapshot_download('zeyadcode/ov_owlv2_model', local_dir='./models/ov_owlv2_model')"    
@@ -32,4 +38,4 @@ RUN chmod +x start.sh
 
 EXPOSE 7860
 
-CMD ["./start.sh"]
+CMD ["./start.sh"]

README.md

@@ -9,6 +9,13 @@ pinned: false
 ---
 
 sudo systemctl stop redis-server
+# Commands
+1. Redis
+> sudo systemctl stop redis-server
+> redis-server --bind 0.0.0.0 --protected-mode no
 
-watchfiles --filter python "celery -A tasks worker -P solo" ./tasks.py
-sudo systemctl stop redis-server && redis-server --bind 0.0.0.0 --protected-mode no
+2. celery worker
+> watchfiles --filter python "cmd /c celery -A tasks worker -P solo" ./tasks.py
+
+3. Run main.py
+> uvicorn main:app --reload

config.py

@@ -1,9 +1,18 @@
+import os
 from pathlib import Path
 
-DEBUG = True
+
+def _env_bool(name: str, default: bool) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+DEBUG = _env_bool("DEBUG", True)
 DEBUG_JOB_ID = "debug-id"
 
-UPLOAD_DIR = Path("uploads")
+UPLOAD_DIR = Path(os.getenv("UPLOAD_DIR", "uploads"))
 INPUT_IMG_NAME = "input_img.png"
 DETECTOR_MODEL_PATH = "./models/ov_owlv2_model/ov_owlv2_model.xml"
 DETECTOR_MODEL_NAME = "google/owlv2-base-patch16-ensemble"
@@ -16,4 +25,12 @@ DETECTOR_OUT_PATH = "detector_boxes.json"
 
 INPAINTOR_OUT_PATH = "inpainted.png"
 
-INPAINTOR_MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
+INPAINTOR_MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
+
+CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL", "redis://localhost:6379/0")
+CELERY_RESULT_BACKEND = os.getenv("CELERY_RESULT_BACKEND", "redis://localhost:6379/1")
+
+JOB_METADATA_NAME = "job.json"
+JOB_TTL_SECONDS = int(os.getenv("JOB_TTL_SECONDS", "3600"))
+JOB_CLEANUP_INTERVAL_SECONDS = int(os.getenv("JOB_CLEANUP_INTERVAL_SECONDS", "1500"))
+RESET_JOBS_ON_STARTUP = _env_bool("RESET_JOBS_ON_STARTUP", not DEBUG)

main.py

@@ -1,3 +1,6 @@
+import asyncio
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
@@ -8,14 +11,28 @@ import tasks as tasks
 from celery.result import AsyncResult
 import config as config
 from schemas import SegmentRequest
-import os
+from services import job_lifecycle
+
+job_lifecycle.ensure_upload_dir()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    if config.RESET_JOBS_ON_STARTUP:
+        job_lifecycle.reset_upload_dir()
 
-if not config.DEBUG:
-    if config.UPLOAD_DIR.exists():
-        shutil.rmtree(config.UPLOAD_DIR)
-    config.UPLOAD_DIR.mkdir(exist_ok=True)
+    cleanup_task = asyncio.create_task(job_lifecycle.cleanup_expired_jobs_forever())
+    try:
+        yield
+    finally:
+        cleanup_task.cancel()
+        try:
+            await cleanup_task
+        except asyncio.CancelledError:
+            pass
 
-app = FastAPI(title="Vandalizer")
+
+app = FastAPI(title="Vandalizer", lifespan=lifespan)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -23,9 +40,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-os.makedirs("uploads", exist_ok=True)
-os.makedirs("outputs", exist_ok=True)  # Good to add this if you use an outputs folder too!
-app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
+app.mount("/uploads", StaticFiles(directory=config.UPLOAD_DIR), name="uploads")
 
 
 @app.get("/")
@@ -36,16 +51,14 @@ def health_check():
 @app.post("/upload")
 def upload(img: UploadFile = File(...)):
     job_id = str(uuid.uuid4())  # uuid1 exposes mac address and time, uuid3,5 uses a key value to generate a hash, same input gets same out, uuid4 is random
-    job_folder = config.UPLOAD_DIR / job_id
-    job_folder.mkdir()
+    job_folder = job_lifecycle.create_job_dir(job_id)
 
     save_path = job_folder / config.INPUT_IMG_NAME
     with open(save_path, "wb") as buffer:
         shutil.copyfileobj(img.file, buffer)
 
-    return job_id
-
-
+    return {"job_id": job_id}
+    
 @app.get("/job_status/{job_id}")
 def get_job_status(job_id: str):
     result = AsyncResult(job_id, app=tasks.celery_app)

services/job_lifecycle.py

@@ -0,0 +1,86 @@
+import asyncio
+import json
+import logging
+import shutil
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+
+import config
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_upload_dir() -> None:
+    config.UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def reset_upload_dir() -> None:
+    if config.UPLOAD_DIR.exists():
+        shutil.rmtree(config.UPLOAD_DIR)
+    ensure_upload_dir()
+
+
+def create_job_dir(job_id: str) -> Path:
+    ensure_upload_dir()
+    job_path = config.UPLOAD_DIR / job_id
+    job_path.mkdir()
+    metadata = {
+        "job_id": job_id,
+        "created_at": datetime.now(timezone.utc).isoformat(),
+    }
+    with open(job_path / config.JOB_METADATA_NAME, "w") as f:
+        json.dump(metadata, f, indent=2)
+    return job_path
+
+
+def _created_at_timestamp(job_path: Path) -> float:
+    metadata_path = job_path / config.JOB_METADATA_NAME
+    if metadata_path.exists():
+        try:
+            with open(metadata_path) as f:
+                metadata = json.load(f)
+            created_at = datetime.fromisoformat(metadata["created_at"])
+            return created_at.timestamp()
+        except (KeyError, ValueError, OSError, json.JSONDecodeError):
+            logger.warning("Ignoring invalid job metadata at %s", metadata_path)
+
+    return job_path.stat().st_mtime
+
+
+def cleanup_expired_jobs(
+    ttl_seconds: int = config.JOB_TTL_SECONDS,
+    now: float | None = None,
+) -> int:
+    ensure_upload_dir()
+    now = time.time() if now is None else now
+    deleted_count = 0
+
+    for job_path in config.UPLOAD_DIR.iterdir():
+        if not job_path.is_dir():
+            continue
+
+        try:
+            age_seconds = now - _created_at_timestamp(job_path)
+        except FileNotFoundError:
+            continue
+
+        if age_seconds < ttl_seconds:
+            continue
+
+        try:
+            shutil.rmtree(job_path)
+            deleted_count += 1
+            logger.info("Deleted expired job folder %s", job_path)
+        except FileNotFoundError:
+            continue
+        except OSError:
+            logger.exception("Failed to delete expired job folder %s", job_path)
+
+    return deleted_count
+
+
+async def cleanup_expired_jobs_forever() -> None:
+    while True:
+        cleanup_expired_jobs()
+        await asyncio.sleep(config.JOB_CLEANUP_INTERVAL_SECONDS)

start.sh

@@ -1,13 +1,24 @@
 #!/bin/bash
 
 # Exit immediately if a command exits with a non-zero status
-set -e
+set -euo pipefail
+
+echo "Resetting runtime job files..."
+mkdir -p uploads outputs /tmp/redis
+rm -rf uploads/* outputs/*
+rm -f dump.rdb /tmp/redis/dump.rdb /var/lib/redis/dump.rdb
 
 echo "Starting Redis server in the background..."
-redis-server --protected-mode no &
+redis-server --protected-mode no --save "" --appendonly no --dir /tmp/redis &
+
+echo "Waiting for Redis..."
+until redis-cli ping >/dev/null 2>&1; do
+    sleep 0.2
+done
+redis-cli FLUSHALL >/dev/null
 
 echo "Starting Celery worker in the background..."
-celery -A tasks worker --loglevel=info &
+celery -A tasks worker --loglevel=info --pool=solo --concurrency=1 &
 
 echo "Starting FastAPI endpoint in the foreground..."
-exec uvicorn main:app --host 0.0.0.0 --port 7860
+exec uvicorn main:app --host 0.0.0.0 --port 7860

tasks.py

@@ -19,8 +19,13 @@ MODELS: dict[str, Any] = {
 
 celery_app = Celery(
     "worker",
-    broker="redis://localhost:6379/0",
-    backend="redis://localhost:6379/1",
+    broker=config.CELERY_BROKER_URL,
+    backend=config.CELERY_RESULT_BACKEND,
+)
+celery_app.conf.update(
+    broker_connection_retry_on_startup=True,
+    result_expires=config.JOB_TTL_SECONDS,
+    task_track_started=True,
 )
 
 
@@ -91,7 +96,7 @@ def segment_objects(self, job_id: str, bboxes=None, points=None, point_labels=No
 
 
 @celery_app.task(bind=True)
-def inpaint(self, job_id: str, prompt: str, negative_prompt: str, num_inference_steps: int, guidance_scale: float) -> bool:
+def inpaint(self, job_id: str, prompt: str, negative_prompt: str = "", num_inference_steps: int = 6, guidance_scale: float = 0.0) -> bool:
     job_path = config.UPLOAD_DIR / job_id
     orig_img = Image.open(job_path / config.INPUT_IMG_NAME).convert("RGB")
     mask_img = Image.open(job_path / config.SEGMENTOR_OUT_BIN_PATH).convert("L")