ClemSummer commited on
Commit
5f629ed
Β·
1 Parent(s): e3a01f4

lazy loading

Browse files
Files changed (1) hide show
  1. main.py +23 -50
main.py CHANGED
@@ -1,55 +1,28 @@
1
- # app/main.py
2
-
3
- from fastapi import FastAPI, UploadFile, File
4
  from fastapi.responses import HTMLResponse
5
- from fastapi.staticfiles import StaticFiles
6
- import shutil
7
- from pathlib import Path
8
  import uvicorn
9
-
10
  from vit_captioning.generate import CaptionGenerator
11
 
12
  app = FastAPI()
13
-
14
- # Serve static files
15
- static_dir = Path(__file__).parent / "vit_captioning" / "static"
16
- app.mount("/static", StaticFiles(directory=static_dir), name="static")
17
-
18
- # βœ… Landing page at `/`
19
- # @app.get("/", response_class=HTMLResponse)
20
- # async def landing():
21
- # return Path("vit_captioning/static/landing.html").read_text()
22
-
23
- @app.get("/")
24
- def read_root():
25
- return {"message": "βœ… Space is working"}
26
-
27
- # βœ… Captioning page at `/captioning`
28
- @app.get("/captioning", response_class=HTMLResponse)
29
- async def captioning():
30
- return Path("vit_captioning/static/captioning/index.html").read_text()
31
-
32
- # βœ… Example: Project 2 placeholder
33
- @app.get("/project2", response_class=HTMLResponse)
34
- async def project2():
35
- return "<h1>Coming Soon: Project 2</h1>"
36
-
37
- # βœ… Caption generation endpoint for captioning app
38
- # Keep the path consistent with your JS fetch()!
39
- caption_generator = CaptionGenerator(
40
- model_type="CLIPEncoder",
41
- checkpoint_path="./vit_captioning/artifacts/CLIPEncoder_40epochs_unfreeze12.pth",
42
- quantized=False
43
- )
44
-
45
- @app.post("/generate")
46
- async def generate(file: UploadFile = File(...)):
47
- temp_file = f"temp_{file.filename}"
48
- with open(temp_file, "wb") as buffer:
49
- shutil.copyfileobj(file.file, buffer)
50
-
51
- captions = caption_generator.generate_caption(temp_file)
52
- return captions
53
-
54
- # if __name__ == "__main__":
55
- # uvicorn.run(app, host="0.0.0.0", port=8000)
 
1
+ from fastapi import FastAPI
 
 
2
  from fastapi.responses import HTMLResponse
 
 
 
3
  import uvicorn
 
4
  from vit_captioning.generate import CaptionGenerator
5
 
6
  app = FastAPI()
7
+ caption_generator = None # Lazy-load placeholder
8
+
9
+ @app.on_event("startup")
10
+ def startup_event():
11
+ global caption_generator
12
+ if caption_generator is None:
13
+ print("Loading CaptionGenerator...")
14
+ caption_generator = CaptionGenerator()
15
+
16
+ @app.get("/", response_class=HTMLResponse)
17
+ def root():
18
+ return "<h3>βœ… Hugging Face Space is alive</h3>"
19
+
20
+ # Example endpoint to trigger model
21
+ @app.get("/caption")
22
+ def caption():
23
+ if caption_generator is None:
24
+ return {"error": "Model not loaded"}
25
+ return {"result": "dummy caption"} # Replace with real logic
26
+
27
+ if __name__ == "__main__":
28
+ uvicorn.run(app, host="0.0.0.0", port=8000)