Update main.py

main.py

@@ -11,11 +11,17 @@ 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
 
-cache_dir = './mycache'
-if not os.path.exists(cache_dir):
-    os.makedirs(cache_dir)
-
-os.environ['TRANSFORMERS_CACHE'] = cache_dir
+# Set cache directory to /tmp/mycache
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/mycache'
+if not os.path.exists('/tmp/mycache'):
+    os.makedirs('/tmp/mycache')
+
+# Initialize the pipeline
+pipe = pipeline(
+    "image-classification", 
+    model="wambugu71/crop_leaf_diseases_vit", 
+    cache_dir='/tmp/mycache'
+)
 
 app = FastAPI()
 
@@ -46,7 +52,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=cache_dir)
+pipe = pipeline("image-classification", model="wambugu71/crop_leaf_diseases_vit", cache_dir="/tmp/mycache")
 @app.post("/api/classify")
 async def classify_image(file: UploadFile = File(...)):
     try: