parallelized caption generation

app.py

@@ -1,4 +1,5 @@
 import os
+import asyncio 
 import torch
 import random
 from PIL import Image
@@ -77,6 +78,31 @@ async def startup_event():
     SEARCH_MODEL = SentenceTransformer('clip-ViT-B-32')
     print("Ensemble is live!")
 
+async def run_inference(m_name, image, temp, top_k, top_p):
+    # This runs in a separate thread to avoid blocking the event loop
+    return await asyncio.to_thread(_generate_sync, m_name, image, temp, top_k, top_p)
+
+def _generate_sync(m_name, image, temp, top_k, top_p):
+    m_data = MODELS[m_name]
+    model = m_data["model"]
+    
+    if m_name == "vit":
+        i_proc, t_proc = m_data["processor"]
+        pixel_values = i_proc(images=image, return_tensors="pt").pixel_values.to(DEVICE)
+        gen_ids = model.generate(
+            pixel_values=pixel_values, max_length=300, do_sample=True, 
+            temperature=temp, top_k=top_k, top_p=top_p
+        )
+        return t_proc.batch_decode(gen_ids, skip_special_tokens=True)[0].strip()
+    else:
+        proc = m_data["processor"]
+        pixel_values = proc(images=image, return_tensors="pt").pixel_values.to(DEVICE)
+        gen_ids = model.generate(
+            pixel_values=pixel_values, max_length=300, do_sample=True, 
+            temperature=temp, top_k=top_k, top_p=top_p
+        )
+        return proc.batch_decode(gen_ids, skip_special_tokens=True)[0].strip()
+
 @app.post("/generate")
 async def generate_endpoint(
     file: UploadFile = File(...),
@@ -85,41 +111,16 @@ async def generate_endpoint(
     top_p: float = Query(0.9)
 ):
     image = Image.open(file.file).convert("RGB")
-    captions = []
-    
-    # Randomly select which models to use for the 5 slots
     available = list(MODELS.keys())
     model_selection = random.choices(available, k=5)
-
-    print("Selected models: ", model_selection)
-
-    for m_name in model_selection:
-        m_data = MODELS[m_name]
-        model = m_data["model"]
-        
-        if m_name == "vit":
-            i_proc, t_proc = m_data["processor"]
-            pixel_values = i_proc(images=image, return_tensors="pt").pixel_values.to(DEVICE)
-            gen_ids = model.generate(
-                pixel_values=pixel_values, max_length=300, do_sample=True, 
-                temperature=temp, top_k=top_k, top_p=top_p
-            )
-            cap = t_proc.batch_decode(gen_ids, skip_special_tokens=True)[0]
-        else:
-            proc = m_data["processor"]
-            pixel_values = proc(images=image, return_tensors="pt").pixel_values.to(DEVICE)
-            gen_ids = model.generate(
-                pixel_values=pixel_values, max_length=300, do_sample=True, 
-                temperature=temp, top_k=top_k, top_p=top_p
-            )
-            cap = proc.batch_decode(gen_ids, skip_special_tokens=True)[0]
-        
-        captions.append(cap.strip())
-
-        print("Caption generated: ", cap.strip())
-
+    
+    # Create tasks for parallel execution
+    tasks = [run_inference(m, image, temp, top_k, top_p) for m in model_selection]
+    captions = await asyncio.gather(*tasks)
+    
     return {"captions": captions, "mix": model_selection}
 
+    
 @app.post("/ui-tester")
 async def ui_tester(file: UploadFile = File(...), description: str = Query(...)):
     """Matches a user description against an image using CLIP embeddings."""