Stanislav commited on
Commit
09d572d
·
1 Parent(s): 7c7b4d8

feat: async lock

Browse files
Files changed (1) hide show
  1. 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 cv2
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
- # 1. Save uploaded image
122
- session_id = uuid.uuid4().hex
123
- save_dir = "uploads"
124
- os.makedirs(save_dir, exist_ok=True)
125
-
126
- image_path = os.path.join(save_dir, f"{session_id}_{image.filename}")
127
- with open(image_path, "wb") as buffer:
128
- shutil.copyfileobj(image.file, buffer)
129
-
130
- # 2. Run main pipeline
131
- process_session_image(
132
- session_id=session_id,
133
- image_path=image_path,
134
- prompt_text=prompt,
135
- sam_wrapper=sam,
136
- dino_wrapper=dino
137
- )
 
138
 
139
- # 3. Redirect to results page
140
- return RedirectResponse(f"/results/{session_id}", status_code=303)
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)