trying to download the model directly

Dockerfile

@@ -16,7 +16,5 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 COPY . .
 
-# Download the 640m model
-RUN wget https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/640m.pt -O /app/640m.pt
 
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "4"]

app.py

@@ -10,6 +10,7 @@ from nudenet import NudeDetector
 import cv2
 from moviepy.editor import VideoFileClip
 from typing import Optional
+import requests
 
 app = FastAPI(
     title="Nudenet API",
@@ -18,10 +19,26 @@ app = FastAPI(
     docs_url="/",
     redoc_url="/redoc"
 )
+def download_model(model_name):
+    url = f"https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/{model_name}"
+    local_path = f"/app/{model_name}"
+    if not os.path.exists(local_path):
+        print(f"Downloading {model_name}...")
+        response = requests.get(url)
+        if response.status_code == 200:
+            with open(local_path, "wb") as f:
+                f.write(response.content)
+            print(f"{model_name} downloaded successfully.")
+        else:
+            print(f"Failed to download {model_name}. Using default model.")
+    return local_path
+
+# Download only the 640m model
+model_640m_path = download_model("640m.pt")
 
 # Initialize NudeDetector with both models
 detector_320n = NudeDetector()
-detector_640m = NudeDetector(model_path="640m.pt", inference_resolution=640)
+detector_640m = NudeDetector(model_path=model_640m_path, inference_resolution=640)
 
 class Base64Image(BaseModel):
     image: str

requirements.txt

@@ -4,4 +4,5 @@ pillow==8.3.2
 nudenet==3.4.2
 opencv-python-headless==4.5.5.64
 moviepy==1.0.3
-python-multipart==0.0.5
+python-multipart==0.0.5
+requests==2.26.0