fastapi added

app.py

@@ -15,6 +15,42 @@ from collections import defaultdict
 from facexlib.utils.misc import download_from_url
 from basicsr.utils.realesrganer import RealESRGANer
 from utils.dataops import auto_split_upscale
+from typing import List, Optional
+
+# FastAPI imports (API server)
+try:
+    from fastapi import FastAPI, UploadFile, File, Form, Depends, HTTPException, status
+    from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+    from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
+except Exception:
+    # Allow Gradio-only usage without FastAPI installed
+    FastAPI = None
+    UploadFile = None
+    File = None
+    Form = None
+    Depends = None
+    HTTPException = Exception
+    status = type("status", (), {"HTTP_401_UNAUTHORIZED": 401, "HTTP_500_INTERNAL_SERVER_ERROR": 500})
+    HTTPBearer = None
+    HTTPAuthorizationCredentials = None
+    FileResponse = None
+    StreamingResponse = None
+    JSONResponse = None
+
+# Mongo imports (optional)
+_mongo_client = None
+_mongo_collection = None
+try:
+    from pymongo import MongoClient
+    _mongo_uri = os.getenv("MONGODB_URI")
+    _mongo_db = os.getenv("MONGODB_DB", "face_upscale")
+    _mongo_col = os.getenv("MONGODB_COLLECTION", "submits")
+    if _mongo_uri:
+        _mongo_client = MongoClient(_mongo_uri, connect=False)
+        _mongo_collection = _mongo_client[_mongo_db][_mongo_col]
+except Exception:
+    _mongo_client = None
+    _mongo_collection = None
 
 input_images_limit = 5
 # Define URLs and their corresponding local storage paths
@@ -1298,4 +1334,158 @@ if __name__ == "__main__":
     parser.add_argument("--input_images_limit", type=int, default=5)
     args = parser.parse_args()
     input_images_limit = args.input_images_limit
-    main()
+    main()
+
+# ---------------------------- FastAPI Application ----------------------------
+
+# Expose FastAPI app named `fastapi_app` so it can be served via: uvicorn app:fastapi_app --host 0.0.0.0 --port 8000
+if FastAPI:
+    fastapi_app = FastAPI(title="Face Upscale API")
+
+    _bearer_scheme = HTTPBearer(auto_error=False) if HTTPBearer else None
+    _api_bearer_token = os.getenv("API_BEARER_TOKEN", "changeme")
+
+    def _verify_bearer_token(credentials: Optional[HTTPAuthorizationCredentials] = Depends(_bearer_scheme) if Depends and _bearer_scheme else None):
+        if not _bearer_scheme:
+            # If FastAPI security is unavailable, deny access
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Unauthorized")
+        if not credentials or credentials.scheme.lower() != "bearer":
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid auth scheme")
+        token = credentials.credentials
+        if not _api_bearer_token or token != _api_bearer_token:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
+        return True
+
+    # Defaults aligned with Gradio UI
+    DEFAULT_FACE_MODEL = 'GFPGANv1.4.pth'
+    DEFAULT_UPSCALE_MODEL = 'SRVGG, realesr-general-x4v3.pth'
+    DEFAULT_SCALE = 4.0
+    DEFAULT_FACE_DET = 'retinaface_resnet50'
+    DEFAULT_FACE_DET_THRESHOLD = 10.0
+    DEFAULT_ONLY_CENTER = False
+    DEFAULT_WITH_MODEL_NAME = True
+
+    # Utility to ensure output dir exists
+    os.makedirs('output', exist_ok=True)
+    os.makedirs('input', exist_ok=True)
+
+    # Singleton Upscale instance for API
+    _api_upscale = Upscale()
+
+    @fastapi_app.get("/health")
+    def health(_: bool = Depends(_verify_bearer_token)):
+        return {"status": "ok"}
+
+    @fastapi_app.get("/source")
+    def source(_: bool = Depends(_verify_bearer_token)):
+        return {
+            "face_models": list(face_models.keys()),
+            "upscale_models": list(typed_upscale_models.keys()),
+            "defaults": {
+                "face_model": DEFAULT_FACE_MODEL,
+                "upscale_model": DEFAULT_UPSCALE_MODEL,
+                "scale": DEFAULT_SCALE,
+                "face_detection": DEFAULT_FACE_DET,
+                "face_detection_threshold": DEFAULT_FACE_DET_THRESHOLD,
+                "face_detection_only_center": DEFAULT_ONLY_CENTER,
+                "with_model_name": DEFAULT_WITH_MODEL_NAME,
+            }
+        }
+
+    @fastapi_app.post("/submit")
+    async def submit(
+        files: List[UploadFile] = File(..., description="One or more image files"),
+        face_model: Optional[str] = Form(DEFAULT_FACE_MODEL),
+        upscale_model: Optional[str] = Form(DEFAULT_UPSCALE_MODEL),
+        scale: float = Form(DEFAULT_SCALE),
+        face_detection: str = Form(DEFAULT_FACE_DET),
+        face_detection_threshold: float = Form(DEFAULT_FACE_DET_THRESHOLD),
+        face_detection_only_center: bool = Form(DEFAULT_ONLY_CENTER),
+        with_model_name: bool = Form(DEFAULT_WITH_MODEL_NAME),
+        _: bool = Depends(_verify_bearer_token)
+    ):
+        try:
+            saved_paths = []
+            for f in files:
+                # Save uploaded file to input directory
+                raw_bytes = await f.read()
+                safe_name = os.path.basename(f.filename)
+                save_path = os.path.join('input', f"{int(time.time()*1000)}_{safe_name}")
+                with open(save_path, 'wb') as out:
+                    out.write(raw_bytes)
+                saved_paths.append(save_path)
+
+            # Build gallery structure expected by Upscale.inference
+            gallery = [[p, ""] for p in saved_paths]
+
+            # Progress stub for API
+            class _NoProgress:
+                def __call__(self, *args, **kwargs):
+                    return None
+            progress_stub = _NoProgress()
+
+            output_files, zip_files = _api_upscale.inference(
+                gallery=gallery,
+                face_restoration=face_model if face_model != "None" else None,
+                upscale_model=upscale_model if upscale_model != "None" else None,
+                scale=float(scale),
+                face_detection=face_detection,
+                face_detection_threshold=face_detection_threshold,
+                face_detection_only_center=face_detection_only_center,
+                outputWithModelName=with_model_name,
+                save_as_png=True,
+                progress=progress_stub,
+            )
+
+            # MongoDB logging
+            if _mongo_collection is not None:
+                try:
+                    _mongo_collection.insert_one({
+                        "ts": int(time.time()),
+                        "files": [os.path.basename(p) for p in saved_paths],
+                        "params": {
+                            "face_model": face_model,
+                            "upscale_model": upscale_model,
+                            "scale": scale,
+                            "face_detection": face_detection,
+                            "face_detection_threshold": face_detection_threshold,
+                            "face_detection_only_center": face_detection_only_center,
+                            "with_model_name": with_model_name,
+                        },
+                        "outputs": {
+                            "images": [os.path.basename(p) for p in (output_files or [])],
+                            "zips": [os.path.basename(p) for p in (zip_files or [])],
+                        },
+                    })
+                except Exception:
+                    pass
+
+            # Build download URLs
+            def _dl_url(name: str) -> str:
+                return f"/download/{name}"
+
+            return {
+                "output_images": [os.path.basename(p) for p in (output_files or [])],
+                "output_zips": [os.path.basename(p) for p in (zip_files or [])],
+                "download_urls": {
+                    "images": [_dl_url(os.path.basename(p)) for p in (output_files or [])],
+                    "zips": [_dl_url(os.path.basename(p)) for p in (zip_files or [])],
+                },
+            }
+        except HTTPException:
+            raise
+        except Exception as e:
+            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
+
+    @fastapi_app.get("/download/{filename}")
+    def download(filename: str, _: bool = Depends(_verify_bearer_token)):
+        safe_name = os.path.basename(filename)
+        # Prefer output dir; fallback to input if requested file is an original
+        candidate_paths = [
+            os.path.join('output', safe_name),
+            os.path.join('input', safe_name),
+        ]
+        for p in candidate_paths:
+            if os.path.isfile(p):
+                return FileResponse(p, filename=safe_name)
+        raise HTTPException(status_code=404, detail="File not found")

requirements.txt

@@ -19,4 +19,8 @@ pyyaml
 yapf
 
 image_gen_aux @ git+https://github.com/huggingface/image_gen_aux
-gdown # supports downloading the large file from Google Drive
+gdown # supports downloading the large file from Google Drive
+fastapi
+uvicorn[standard]
+pymongo
+python-multipart