Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 import numpy as np
-from briarmbg import BriaRMBG
+from architecture import BackgroundEnhancer
 import gradio as gr
 import os
 
@@ -34,7 +34,7 @@ def example_inference(image):
 
     model_path = "model.pth"
 
-    net = BriaRMBG()
+    net = BackgroundEnhancer()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     net.load_state_dict(torch.load(model_path, map_location=device))
     net.to(device)