Update main.py

main.py

@@ -11,7 +11,11 @@ from projects.DL_CatDog.DL_CatDog import preprocess_image, read_image, model_DL_
 from projects.ML_StudentPerformance.ML_StudentPerformace import predict_student_performance, create_custom_data, form1
 from projects.ML_DiabetesPrediction.ML_DiabetesPrediction import model_ML_DiabetesPrediction, form2
 
-os.environ['HF_HOME'] = './mycache'
+cache_dir = './mycache'
+if not os.path.exists(cache_dir):
+    os.makedirs(cache_dir)
+
+os.environ['TRANSFORMERS_CACHE'] = cache_dir
 
 app = FastAPI()
 
@@ -24,7 +28,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
 # Health check route
 @app.get("/api/working")
 def home():
@@ -43,7 +46,7 @@ async def predict_DL_CatDog(file: UploadFile = File(...)):
         return JSONResponse(content={"ok": -1, "message": f"Something went wrong! {str(e)}"}, status_code=500)
 
 # Classification route for DL_PlantDisease
-pipe = pipeline("image-classification", model="wambugu71/crop_leaf_diseases_vit", cache_dir="./mycache")
+pipe = pipeline("image-classification", model="wambugu71/crop_leaf_diseases_vit", cache_dir=cache_dir)
 @app.post("/api/classify")
 async def classify_image(file: UploadFile = File(...)):
     try: