Codewithsalty commited on
Commit
bc39385
·
verified ·
1 Parent(s): daee81e

Update newapi.py

Browse files
Files changed (1) hide show
  1. newapi.py +5 -5
newapi.py CHANGED
@@ -11,13 +11,13 @@ from huggingface_hub import hf_hub_download
11
  from models.TumorModel import TumorClassification, GliomaStageModel
12
  from utils import get_precautions_from_gemini
13
 
14
- # Create writable cache directory for Hugging Face
15
- os.makedirs("./cache", exist_ok=True)
16
 
17
  # Initialize FastAPI app
18
  app = FastAPI(title="Brain Tumor Detection API")
19
 
20
- # Enable CORS for all origins (adjust for production)
21
  app.add_middleware(
22
  CORSMiddleware,
23
  allow_origins=["*"],
@@ -30,7 +30,7 @@ app.add_middleware(
30
  btd_model_path = hf_hub_download(
31
  repo_id="Codewithsalty/brain-tumor-models",
32
  filename="BTD_model.pth",
33
- cache_dir="./cache" # ✅ fix
34
  )
35
  tumor_model = TumorClassification()
36
  tumor_model.load_state_dict(torch.load(btd_model_path, map_location="cpu"))
@@ -40,7 +40,7 @@ tumor_model.eval()
40
  glioma_model_path = hf_hub_download(
41
  repo_id="Codewithsalty/brain-tumor-models",
42
  filename="glioma_stages.pth",
43
- cache_dir="./cache" # ✅ fix
44
  )
45
  glioma_model = GliomaStageModel()
46
  glioma_model.load_state_dict(torch.load(glioma_model_path, map_location="cpu"))
 
11
  from models.TumorModel import TumorClassification, GliomaStageModel
12
  from utils import get_precautions_from_gemini
13
 
14
+ # Use Hugging Face's writable cache directory
15
+ cache_dir = os.getenv("HF_HOME", "/data")
16
 
17
  # Initialize FastAPI app
18
  app = FastAPI(title="Brain Tumor Detection API")
19
 
20
+ # Enable CORS for all origins
21
  app.add_middleware(
22
  CORSMiddleware,
23
  allow_origins=["*"],
 
30
  btd_model_path = hf_hub_download(
31
  repo_id="Codewithsalty/brain-tumor-models",
32
  filename="BTD_model.pth",
33
+ cache_dir=cache_dir
34
  )
35
  tumor_model = TumorClassification()
36
  tumor_model.load_state_dict(torch.load(btd_model_path, map_location="cpu"))
 
40
  glioma_model_path = hf_hub_download(
41
  repo_id="Codewithsalty/brain-tumor-models",
42
  filename="glioma_stages.pth",
43
+ cache_dir=cache_dir
44
  )
45
  glioma_model = GliomaStageModel()
46
  glioma_model.load_state_dict(torch.load(glioma_model_path, map_location="cpu"))