Create main.py

main.py

@@ -0,0 +1,46 @@
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
+import io
+import torch
+from clip_interrogator import Config, Interrogator
+
+app = FastAPI()
+
+# Allow CORS for all origins (adjust as needed for production)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Setup the CLIP Interrogator
+config = Config()
+config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+config.blip_offload = False if torch.cuda.is_available() else True
+config.chunk_size = 2048
+config.flavor_intermediate_count = 512
+config.blip_num_beams = 64
+
+ci = Interrogator(config)
+
+@app.post("/inference/")
+async def interrogate_images(file: UploadFile = File(...), mode: str = Form(...), best_max_flavors: int = Form(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert('RGB')
+        
+        if mode == 'best':
+            prompt_result = ci.interrogate(image, max_flavors=int(best_max_flavors))
+        elif mode == 'classic':
+            prompt_result = ci.interrogate_classic(image)
+        else:
+            prompt_result = ci.interrogate_fast(image)
+        
+        return JSONResponse(content={"prompt_results": [prompt_result]})
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+