fixed convnextv2 bug

Dockerfile

@@ -1,15 +1,14 @@
 FROM python:3.9-slim
 
-ENV TRANSFORMERS_CACHE=/data/.cache
+ENV TRANSFORMERS_CACHE=/data/.cache/transformers
+ENV HF_HOME=/data/.cache/huggingface
+ENV MPLCONFIGDIR=/data/.cache/matplotlib
 
 WORKDIR /code
 
 COPY requirements.txt .
-
 RUN pip install --no-cache-dir -r requirements.txt
-
 COPY . .
 
 EXPOSE 7860
-
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,20 +1,25 @@
+import os
+
+os.environ['TRANSFORMERS_CACHE'] = '/data/.cache/transformers'
+os.environ['HF_HOME'] = '/data/.cache/huggingface'
+os.environ['MPLCONFIGDIR'] = '/data/.cache/matplotlib'
+
 import torch
 import torch.nn as nn
 import yaml
 from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
-import os
 from transformers import ConvNextV2ForImageClassification
 from typing import Dict, Tuple
 
 MODEL_CHECKPOINTS = {
-    "ConvNeXt Tiny (Best)": "checkpoints/convnext_v2_tiny_best.pth",
+    "ConvNeXt tiny (Best)": "checkpoints/convnext_v2_tiny_best.pth",
     "EfficientNet-B0": "checkpoints/effnet_b0_best.pth",
     "EfficientNet-B3": "checkpoints/effnet_b3_best.pth",
     "Vision Transformer B-16": "checkpoints/vit_b_16_best.pth"
 }
-DEFAULT_MODEL_NAME = "ConvNeXt Tiny (Best)"
+DEFAULT_MODEL_NAME = "ConvNeXt tiny (Best)"
 
 MODELS: Dict[str, Tuple[nn.Module, Dict[int, str]]] = {}
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -48,7 +53,6 @@ def get_model(model_name: str, num_classes: int) -> nn.Module:
     return model
 
 def load_checkpoint(checkpoint_path: str, device: torch.device) -> Tuple[nn.Module, Dict[int, str]]:
-    # (Same load_checkpoint function as in main.py)
     if not os.path.exists(checkpoint_path):
         raise FileNotFoundError(f"Checkpoint file not found at: {checkpoint_path}")
     checkpoint = torch.load(checkpoint_path, map_location=device)

main.py

@@ -1,10 +1,15 @@
+import os
+
+os.environ['TRANSFORMERS_CACHE'] = '/data/.cache/transformers'
+os.environ['HF_HOME'] = '/data/.cache/huggingface'
+os.environ['MPLCONFIGDIR'] = '/data/.cache/matplotlib'
+
 import torch
 import torch.nn as nn
 import yaml
 from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
-import os
 import base64
 import io
 import time
@@ -16,14 +21,13 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import ConvNextV2ForImageClassification
 
-
 MODEL_CHECKPOINTS = {
-    "ConvNeXt Tiny (Best)": "checkpoints/convnext_v2_tiny_best.pth",
+    "ConvNeXt tiny (Best)": "checkpoints/convnext_v2_tiny_best.pth",
     "EfficientNet-B0": "checkpoints/effnet_b0_best.pth",
     "EfficientNet-B3": "checkpoints/effnet_b3_best.pth",
     "Vision Transformer B-16": "checkpoints/vit_b_16_best.pth"
 }
-DEFAULT_MODEL_NAME = "ConvNeXt Tiny (Best)"
+DEFAULT_MODEL_NAME = "ConvNeXt tiny (Best)"
 
 GPU_MODELS: Dict[str, Tuple[nn.Module, Dict[int, str]]] = {}
 CPU_MODELS: Dict[str, Tuple[nn.Module, Dict[int, str]]] = {}
@@ -79,11 +83,9 @@ gpu_device = torch.device("cuda") if torch.cuda.is_available() else None
 for display_name, ckpt_path in MODEL_CHECKPOINTS.items():
     if os.path.exists(ckpt_path):
         print(f"Loading '{display_name}'...")
-        # Load for CPU (always)
         cpu_model, idx_to_class = load_checkpoint(ckpt_path, cpu_device)
         CPU_MODELS[display_name] = (cpu_model, idx_to_class)
         print(f"  > Loaded '{display_name}' for CPU.")
-        # Load for GPU if available
         if gpu_device:
             gpu_model, _ = load_checkpoint(ckpt_path, gpu_device)
             GPU_MODELS[display_name] = (gpu_model, idx_to_class)