Refactor for Hugging Face Spaces: Remove Redis/Celery, add standalone mode

Dockerfile

@@ -13,5 +13,11 @@ COPY . /app
 ENV PYTHONUNBUFFERED=1
 ENV PYTHONDONTWRITEBYTECODE=1
 
-# For GPU: use an nvidia/cuda base image and install the correct torch wheel with CUDA.
-# When running with nvidia runtime, pass --gpus=all to docker run or set deploy settings in compose.
+# Create directories with permissions
+RUN mkdir -p /tmp/uploads /tmp/results && chmod 777 /tmp/uploads /tmp/results
+
+# Expose Hugging Face default port
+EXPOSE 7860
+
+# Start the application
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,55 +1,39 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
 from fastapi.responses import FileResponse, JSONResponse
 from uuid import uuid4
 from pathlib import Path
 import shutil
 import os
 import json
-import redis
-from celery import Celery
 from dotenv import load_dotenv
+from tasks import process_image_task
 
 load_dotenv()
 
-# Directories (mounted by docker-compose)
-UPLOAD_DIR = Path(os. environ. get("UPLOAD_DIR", "/data/uploads"))
-RESULT_DIR = Path(os.environ. get("RESULT_DIR", "/data/results"))
+# Directories
+# Use /tmp for Hugging Face Spaces as it is writable
+UPLOAD_DIR = Path("/tmp/uploads")
+RESULT_DIR = Path("/tmp/results")
 UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
 RESULT_DIR.mkdir(parents=True, exist_ok=True)
 
-# Redis for job status
-REDIS_HOST = os. environ.get("REDIS_HOST", "redis")
-REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))
-CELERY_BROKER_URL = os.environ.get("CELERY_BROKER_URL", "redis://redis:6379/0")
+# In-memory job store (Global variable)
+# Since HF Spaces (Free) runs 1 replica, this works for a demo.
+JOBS = {}
 
-# Lazy connections
-_rdb = None
-_celery_client = None
-
-def get_redis():
-    global _rdb
-    if _rdb is None:
-        _rdb = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
-    return _rdb
-
-def get_celery():
-    global _celery_client
-    if _celery_client is None:
-        _celery_client = Celery(broker=CELERY_BROKER_URL)
-    return _celery_client
-
-app = FastAPI(title="Depth->STL processing service (API)")
-
-
-def set_status(job_id: str, state: str, detail: str = "", result: str = ""):
-    payload = {"state": state, "detail": detail, "result": result}
-    get_redis().set(job_id, json.dumps(payload))
+app = FastAPI(title="Depth->STL processing service (Standalone)")
 
+def update_job_status(job_id: str, state: str, detail: str = "", result: str = ""):
+    JOBS[job_id] = {
+        "state": state,
+        "detail": detail,
+        "result": result
+    }
 
 @app.post("/upload/")
-async def upload_image(file: UploadFile = File(... )):
+async def upload_image(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
     # Basic validation
-    if not file.content_type. startswith("image/"):
+    if not file.content_type.startswith("image/"):
         raise HTTPException(status_code=400, detail="File must be an image")
 
     job_id = str(uuid4())
@@ -57,46 +41,47 @@ async def upload_image(file: UploadFile = File(... )):
     fname = f"{job_id}_{safe_name}"
     save_path = UPLOAD_DIR / fname
 
-    # Save uploaded file to mounted volume
+    # Save uploaded file
     try:
         with save_path.open("wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed to save upload: {e}")
 
-    # Mark queued and enqueue Celery task
-    set_status(job_id, "QUEUED", "Job received and queued")
-    try:
-        async_result = get_celery().send_task(
-            "tasks.process_image_task",
-            args=[str(save_path), str(RESULT_DIR), job_id],
-            kwargs={},
-            queue=os.environ.get("CELERY_QUEUE", None),
-        )
-    except Exception as e:
-        set_status(job_id, "FAILURE", f"Failed to enqueue task: {e}")
-        raise HTTPException(status_code=500, detail=f"Failed to enqueue task: {e}")
+    # Mark queued
+    update_job_status(job_id, "QUEUED", "Job received and queued")
+
+    # Add to background tasks
+    background_tasks.add_task(
+        process_image_task,
+        str(save_path),
+        str(RESULT_DIR),
+        job_id,
+        update_job_status
+    )
 
-    return {"job_id": job_id, "celery_id": str(async_result.id)}
+    return {"job_id": job_id}
 
 
 @app.get("/status/{job_id}")
 def status(job_id: str):
-    raw = get_redis().get(job_id)
-    if not raw:
+    job = JOBS.get(job_id)
+    if not job:
         return JSONResponse({"state": "UNKNOWN", "detail": "No such job_id"}, status_code=404)
-    return JSONResponse(json.loads(raw))
+    return JSONResponse(job)
 
 
 @app.get("/download/{job_id}")
 def download(job_id: str):
-    raw = get_redis().get(job_id)
-    if not raw:
+    job = JOBS.get(job_id)
+    if not job:
         raise HTTPException(status_code=404, detail="No such job")
-    info = json.loads(raw)
-    if info.get("state") != "SUCCESS":
+    
+    if job.get("state") != "SUCCESS":
         raise HTTPException(status_code=404, detail="Result not ready")
-    stl_path = info.get("result")
+    
+    stl_path = job.get("result")
     if not stl_path or not Path(stl_path).exists():
         raise HTTPException(status_code=404, detail="Result file missing")
+    
     return FileResponse(path=stl_path, filename=Path(stl_path).name, media_type="application/sla")

tasks.py

@@ -1,9 +1,7 @@
 import os
-from celery import Celery
 from pathlib import Path
 import traceback
 import json
-import redis
 import time
 import sys
 
@@ -11,24 +9,17 @@ import numpy as np
 import cv2
 import open3d as o3d
 import torch
-from PIL import Image
 import trimesh
 from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 
-# Celery / Redis config
-CELERY_BROKER = os.environ.get("CELERY_BROKER_URL", "redis://redis:6379/0")
-CELERY_BACKEND = os.environ.get("CELERY_RESULT_BACKEND", "redis://redis:6379/0")
-REDIS_HOST = os.environ.get("REDIS_HOST", "redis")
-REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))
-
-# Pipeline settings (fixed to orthographic + Poisson depth=9 to match your notebook)
-DEPTH_CHECKPOINT = os.environ.get("DEPTH_CHECKPOINT", "/models/depth-anything-Large-hf")
-USE_GPU = int(os.environ.get("USE_GPU", "1"))
+# Pipeline settings
+DEPTH_CHECKPOINT = os.environ.get("DEPTH_CHECKPOINT", "LiheYoung/depth-anything-large-hf") # Default to HF Hub model if local not found
+USE_GPU = int(os.environ.get("USE_GPU", "0")) # Default to CPU for HF Spaces
 POISSON_DEPTH = int(os.environ.get("POISSON_DEPTH", "9"))
 OUTLIER_NEIGHBORS = int(os.environ.get("OUTLIER_NEIGHBORS", "15"))
 OUTLIER_STD_RATIO = float(os.environ.get("OUTLIER_STD_RATIO", "1.0"))
-ORTHO_SCALE_FACTOR = float(os.environ.get("ORTHO_SCALE_FACTOR", "255"))  # same as your function
-INFERENCE_RESIZE = int(os.environ.get("INFERENCE_RESIZE", "0"))  # 0 keeps original
+ORTHO_SCALE_FACTOR = float(os.environ.get("ORTHO_SCALE_FACTOR", "255"))
+INFERENCE_RESIZE = int(os.environ.get("INFERENCE_RESIZE", "0"))
 RESULT_PREFIX = os.environ.get("RESULT_PREFIX", "")
 
 try:
@@ -36,9 +27,6 @@ try:
 except Exception:
     pass
 
-celery = Celery("tasks", broker=CELERY_BROKER, backend=CELERY_BACKEND)
-rdb = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
-
 _model = None
 _processor = None
 _device = "cpu"
@@ -47,10 +35,6 @@ def log(msg):
     print(msg, flush=True)
     sys.stdout.flush()
 
-def set_status(job_id: str, state: str, detail: str = "", result: str = ""):
-    payload = {"state": state, "detail": detail, "result": result}
-    rdb.set(job_id, json.dumps(payload))
-
 def load_model():
     global _model, _processor, _device
     if _model is None:
@@ -85,11 +69,10 @@ def build_orthographic_point_cloud(depth_u8: np.ndarray, color_rgb: np.ndarray)
     pcd.colors = o3d.utility.Vector3dVector(colors)
     return pcd
 
-@celery.task(bind=True)
-def process_image_task(self, image_path: str, result_dir: str, job_id: str):
+def process_image_task(image_path: str, result_dir: str, job_id: str, status_callback):
     start = time.time()
     try:
-        set_status(job_id, "RUNNING", "Loading model")
+        status_callback(job_id, "RUNNING", "Loading model")
         model, processor, device = load_model()
         log(f"[{job_id}] Model loaded on {device}")
 
@@ -108,7 +91,7 @@ def process_image_task(self, image_path: str, result_dir: str, job_id: str):
         else:
             img_proc = img_rgb
 
-        set_status(job_id, "RUNNING", "Running depth inference")
+        status_callback(job_id, "RUNNING", "Running depth inference")
         depth_inputs = processor(images=img_proc, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = model(**depth_inputs)
@@ -120,7 +103,7 @@ def process_image_task(self, image_path: str, result_dir: str, job_id: str):
 
         depth_u8 = normalize_depth_uint8(depth)
 
-        set_status(job_id, "RUNNING", "Building orthographic point cloud")
+        status_callback(job_id, "RUNNING", "Building orthographic point cloud")
         pcd = build_orthographic_point_cloud(depth_u8, color_resized)
 
         # Outlier removal (nb=15, std_ratio=1.0)
@@ -144,7 +127,7 @@ def process_image_task(self, image_path: str, result_dir: str, job_id: str):
         if num_pts == 0:
             raise RuntimeError("Empty point cloud after cleanup")
 
-        set_status(job_id, "RUNNING", f"Poisson reconstruction depth={POISSON_DEPTH}")
+        status_callback(job_id, "RUNNING", f"Poisson reconstruction depth={POISSON_DEPTH}")
         mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
             pcd, depth=POISSON_DEPTH
         )
@@ -165,14 +148,14 @@ def process_image_task(self, image_path: str, result_dir: str, job_id: str):
         Path(result_dir).mkdir(parents=True, exist_ok=True)
         stl_path = Path(result_dir) / f"{RESULT_PREFIX}{job_id}.stl"
 
-        set_status(job_id, "RUNNING", "Exporting STL")
+        status_callback(job_id, "RUNNING", "Exporting STL")
         tm = trimesh.Trimesh(vertices=np.asarray(mesh.vertices),
                              faces=np.asarray(mesh.triangles),
                              process=True)
         tm.export(str(stl_path), file_type="stl")
 
         total = time.time() - start
-        set_status(job_id, "SUCCESS", f"Done in {total:.2f}s", str(stl_path))
+        status_callback(job_id, "SUCCESS", f"Done in {total:.2f}s", str(stl_path))
         log(f"[{job_id}] SUCCESS total={total:.2f}s STL={stl_path}")
         return {
             "status": "success",
@@ -181,6 +164,6 @@ def process_image_task(self, image_path: str, result_dir: str, job_id: str):
         }
     except Exception as e:
         traceback.print_exc()
-        set_status(job_id, "FAILURE", str(e))
+        status_callback(job_id, "FAILURE", str(e))
         log(f"[{job_id}] FAILURE: {e}")
         raise