lock with async wait

app.py

@@ -17,6 +17,8 @@ import uuid
 import logging
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from asyncio import Lock
+
 
 logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
 
@@ -24,7 +26,7 @@ MAX_SEED = np.iinfo(np.int32).max
 USE_TORCH_COMPILE = os.environ.get("USE_TORCH_COMPILE", "0") == "1"
 SPACE_ID = os.environ.get("SPACE_ID", "")
 DEV = os.environ.get("DEV", "0") == "1"
-os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 DB_PATH = Path("/data/cache") if SPACE_ID else Path("./cache")
 IMGS_PATH = DB_PATH / "imgs"
@@ -32,6 +34,7 @@ DB_PATH.mkdir(exist_ok=True, parents=True)
 IMGS_PATH.mkdir(exist_ok=True, parents=True)
 
 database = Database(DB_PATH)
+generate_lock = Lock()
 
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -120,21 +123,25 @@ app.add_middleware(
 
 
 @app.get("/image")
-async def generate_image(prompt: str, negative_prompt: str = "", seed: int = 2134213213):
+async def generate_image(
+    prompt: str, negative_prompt: str = "", seed: int = 2134213213
+):
     cached_img = database.check(prompt, negative_prompt, seed)
     if cached_img:
         logging.info(f"Image found in cache: {cached_img[0]}")
         return StreamingResponse(open(cached_img[0], "rb"), media_type="image/jpeg")
 
     logging.info(f"Image not found in cache, generating new image")
-    pil_image = generate(prompt, negative_prompt, seed)
-    img_id = str(uuid.uuid4())
-    img_path = IMGS_PATH / f"{img_id}.jpg"
-    pil_image.save(img_path)
-    img_io = io.BytesIO()
-    pil_image.save(img_io, "JPEG")
-    img_io.seek(0)
-    database.insert(prompt, negative_prompt, str(img_path), seed)
+    async with generate_lock:
+
+        pil_image = generate(prompt, negative_prompt, seed)
+        img_id = str(uuid.uuid4())
+        img_path = IMGS_PATH / f"{img_id}.jpg"
+        pil_image.save(img_path)
+        img_io = io.BytesIO()
+        pil_image.save(img_io, "JPEG")
+        img_io.seek(0)
+        database.insert(prompt, negative_prompt, str(img_path), seed)
 
     return StreamingResponse(img_io, media_type="image/jpeg")