Spaces:
Sleeping
Sleeping
Stanislav
commited on
Commit
·
09d572d
1
Parent(s):
7c7b4d8
feat: async lock
Browse files- run_fastapi.py +23 -25
run_fastapi.py
CHANGED
@@ -5,7 +5,7 @@ from fastapi.templating import Jinja2Templates
|
|
5 |
import uuid
|
6 |
import shutil
|
7 |
import os
|
8 |
-
import
|
9 |
|
10 |
from pipeline.process_session import process_session_image
|
11 |
from database.db import init_db
|
@@ -15,6 +15,7 @@ from models.dino import DinoWrapper
|
|
15 |
|
16 |
from huggingface_hub import hf_hub_download
|
17 |
|
|
|
18 |
|
19 |
app = FastAPI()
|
20 |
|
@@ -32,13 +33,9 @@ print("WRITE to /tmp:", os.access("/tmp", os.W_OK))
|
|
32 |
BASE_DIR = "/tmp" if os.access("/tmp", os.W_OK) else "."
|
33 |
|
34 |
WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
|
35 |
-
UPLOADS_DIR = os.path.join(BASE_DIR, "uploads")
|
36 |
-
OUTPUTS_DIR = os.path.join(BASE_DIR, "outputs")
|
37 |
|
38 |
# Create directories if not exist
|
39 |
os.makedirs(WEIGHTS_DIR, exist_ok=True)
|
40 |
-
os.makedirs(UPLOADS_DIR, exist_ok=True)
|
41 |
-
os.makedirs(OUTPUTS_DIR, exist_ok=True)
|
42 |
|
43 |
# --- Initialize database
|
44 |
init_db()
|
@@ -117,27 +114,28 @@ def show_results(request: Request, session_id: str):
|
|
117 |
|
118 |
|
119 |
@app.post("/process")
|
120 |
-
def process_image(request: Request, image: UploadFile = Form(...), prompt: str = Form(...)):
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
|
|
138 |
|
139 |
-
|
140 |
-
|
141 |
|
142 |
|
143 |
@app.post("/finalize", response_class=HTMLResponse)
|
|
|
5 |
import uuid
|
6 |
import shutil
|
7 |
import os
|
8 |
+
import asyncio
|
9 |
|
10 |
from pipeline.process_session import process_session_image
|
11 |
from database.db import init_db
|
|
|
15 |
|
16 |
from huggingface_hub import hf_hub_download
|
17 |
|
18 |
+
process_lock = asyncio.Lock()
|
19 |
|
20 |
app = FastAPI()
|
21 |
|
|
|
33 |
BASE_DIR = "/tmp" if os.access("/tmp", os.W_OK) else "."
|
34 |
|
35 |
WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
|
|
|
|
|
36 |
|
37 |
# Create directories if not exist
|
38 |
os.makedirs(WEIGHTS_DIR, exist_ok=True)
|
|
|
|
|
39 |
|
40 |
# --- Initialize database
|
41 |
init_db()
|
|
|
114 |
|
115 |
|
116 |
@app.post("/process")
|
117 |
+
async def process_image(request: Request, image: UploadFile = Form(...), prompt: str = Form(...)):
|
118 |
+
async with process_lock:
|
119 |
+
# 1. Save uploaded image
|
120 |
+
session_id = uuid.uuid4().hex
|
121 |
+
save_dir = "uploads"
|
122 |
+
os.makedirs(save_dir, exist_ok=True)
|
123 |
+
|
124 |
+
image_path = os.path.join(save_dir, f"{session_id}_{image.filename}")
|
125 |
+
with open(image_path, "wb") as buffer:
|
126 |
+
shutil.copyfileobj(image.file, buffer)
|
127 |
+
|
128 |
+
# 2. Run main pipeline
|
129 |
+
process_session_image(
|
130 |
+
session_id=session_id,
|
131 |
+
image_path=image_path,
|
132 |
+
prompt_text=prompt,
|
133 |
+
sam_wrapper=sam,
|
134 |
+
dino_wrapper=dino
|
135 |
+
)
|
136 |
|
137 |
+
# 3. Redirect to results page
|
138 |
+
return RedirectResponse(f"/results/{session_id}", status_code=303)
|
139 |
|
140 |
|
141 |
@app.post("/finalize", response_class=HTMLResponse)
|