Spaces:
Running
Running
Commit ·
295de8e
1
Parent(s): 74d6a33
fix: added
Browse files- app/main.py +7 -2
app/main.py
CHANGED
|
@@ -11,6 +11,7 @@ from transformers import BlipForConditionalGeneration, BlipProcessor
|
|
| 11 |
Image.MAX_IMAGE_PIXELS = None
|
| 12 |
|
| 13 |
MODEL_ID = os.getenv("MODEL_ID", "Salesforce/blip-image-captioning-large")
|
|
|
|
| 14 |
MAX_IMAGE_SIZE = (1024, 1024)
|
| 15 |
|
| 16 |
processor: BlipProcessor | None = None
|
|
@@ -20,7 +21,7 @@ model: BlipForConditionalGeneration | None = None
|
|
| 20 |
def load_model() -> None:
|
| 21 |
global processor, model
|
| 22 |
if processor is None:
|
| 23 |
-
processor = BlipProcessor.from_pretrained(MODEL_ID)
|
| 24 |
if model is None:
|
| 25 |
model = BlipForConditionalGeneration.from_pretrained(MODEL_ID)
|
| 26 |
|
|
@@ -70,7 +71,11 @@ async def root() -> dict[str, str]:
|
|
| 70 |
|
| 71 |
@app.get("/health")
|
| 72 |
async def health() -> dict[str, str]:
|
| 73 |
-
return {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 74 |
|
| 75 |
|
| 76 |
@app.post("/caption")
|
|
|
|
| 11 |
Image.MAX_IMAGE_PIXELS = None
|
| 12 |
|
| 13 |
MODEL_ID = os.getenv("MODEL_ID", "Salesforce/blip-image-captioning-large")
|
| 14 |
+
USE_FAST_PROCESSOR = True
|
| 15 |
MAX_IMAGE_SIZE = (1024, 1024)
|
| 16 |
|
| 17 |
processor: BlipProcessor | None = None
|
|
|
|
| 21 |
def load_model() -> None:
|
| 22 |
global processor, model
|
| 23 |
if processor is None:
|
| 24 |
+
processor = BlipProcessor.from_pretrained(MODEL_ID, use_fast=USE_FAST_PROCESSOR)
|
| 25 |
if model is None:
|
| 26 |
model = BlipForConditionalGeneration.from_pretrained(MODEL_ID)
|
| 27 |
|
|
|
|
| 71 |
|
| 72 |
@app.get("/health")
|
| 73 |
async def health() -> dict[str, str]:
|
| 74 |
+
return {
|
| 75 |
+
"status": "ok",
|
| 76 |
+
"model_id": MODEL_ID,
|
| 77 |
+
"use_fast_processor": str(USE_FAST_PROCESSOR).lower(),
|
| 78 |
+
}
|
| 79 |
|
| 80 |
|
| 81 |
@app.post("/caption")
|