Model changed and weights

app.py

@@ -1,10 +1,10 @@
 import gradio as gr 
-from PIL import Image
 import torch
 import torchvision.transforms as transforms
 import numpy as np
 
-from archs.model import FourNet
+from PIL import Image
+from model.flol import create_model
 
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -13,11 +13,11 @@ pil_to_tensor = transforms.ToTensor()
     
 # define some parameters based on the run we want to make
 
-model = FourNet(nf = 16)
+model = create_model()
 
-checkpoints = torch.load('./models/NAFourNet16_LOLv2Real.pt', map_location=device)
+checkpoints = torch.load('./weights/flolv2_UHDLL.pt', map_location=device)
 
-model.load_state_dict(checkpoints['model_state_dict'])
+model.load_state_dict(checkpoints['params'])
 
 model = model.to(device)