khushalcodiste commited on
Commit
295de8e
·
1 Parent(s): 74d6a33

fix: added

Browse files
Files changed (1) hide show
  1. 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 {"status": "ok", "model_id": MODEL_ID}
 
 
 
 
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")