DonImages commited on
Commit
30fa026
·
verified ·
1 Parent(s): 0cefee0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -1,14 +1,16 @@
1
  from fastapi import FastAPI, HTTPException
2
  import base64
3
  import os
 
4
 
5
  app = FastAPI()
6
 
7
  # Global variable to hold the LoRA weights
8
  lora_weights = None
9
 
10
- @app.on_event("startup")
11
- async def load_lora_weights():
 
12
  global lora_weights
13
  lora_path = "./lora_file.pth" # Ensure the correct file name
14
  if os.path.exists(lora_path):
@@ -17,6 +19,10 @@ async def load_lora_weights():
17
  print("LoRA weights loaded and preprocessed successfully.")
18
  else:
19
  raise HTTPException(status_code=500, detail="LoRA file not found.")
 
 
 
 
20
 
21
  @app.post("/modify-prompt")
22
  async def modify_prompt(prompt: str):
 
1
  from fastapi import FastAPI, HTTPException
2
  import base64
3
  import os
4
+ from contextlib import asynccontextmanager
5
 
6
  app = FastAPI()
7
 
8
  # Global variable to hold the LoRA weights
9
  lora_weights = None
10
 
11
+ # Lifespan event to load the LoRA weights
12
+ @asynccontextmanager
13
+ async def lifespan(app: FastAPI):
14
  global lora_weights
15
  lora_path = "./lora_file.pth" # Ensure the correct file name
16
  if os.path.exists(lora_path):
 
19
  print("LoRA weights loaded and preprocessed successfully.")
20
  else:
21
  raise HTTPException(status_code=500, detail="LoRA file not found.")
22
+ yield
23
+ # Cleanup if necessary (but not required in this case)
24
+
25
+ app.add_event_handler("lifespan", lifespan)
26
 
27
  @app.post("/modify-prompt")
28
  async def modify_prompt(prompt: str):