Update app.py

app.py

@@ -1,12 +1,12 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-from torchvision.transforms.functional import normalize  # 수정된 부분
+from torchvision.transforms.functional import normalize
 from huggingface_hub import hf_hub_download
 from briarmbg import BriaRMBG
 from PIL import Image
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, JSONResponse
 import os
 
 app = FastAPI()
@@ -15,10 +15,10 @@ app = FastAPI()
 net = BriaRMBG()
 model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
 if torch.cuda.is_available():
-    net.load_state_dict(torch.load(model_path))
+    net.load_state_dict(torch.load(model_path, map_location="cuda", weights_only=True))
     net = net.cuda()
 else:
-    net.load_state_dict(torch.load(model_path, map_location="cpu"))
+    net.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=True))
 net.eval()
 
 def resize_image(image):
@@ -58,6 +58,10 @@ def process_image(image: Image.Image):
     
     return output_path
 
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to the Background Removal API"}
+
 @app.post("/remove-background/")
 async def remove_background(file: UploadFile = File(...)):
     image = Image.open(file.file)