Update app.py

app.py

@@ -7,6 +7,7 @@ import argparse
 import shutil
 import gradio as gr
 from PIL import Image
+from runpy import run_path
 
 examples = [['sample1.png'],
             ['sample2.png']]
@@ -40,12 +41,13 @@ def inference(img, task, run_on):
         height = int(scale*height)
         img = img.resize((width,height), Image.ANTIALIAS)
       
-    #parameters = {'inp_channels':3, 'out_channels':3, 'dim':48, 'num_blocks':[4,6,6,8], 'num_refinement_blocks':4, 'heads':[1,2,4,8], 'ffn_expansion_factor':2.66, 'bias':False, 'LayerNorm_type':'WithBias', 'dual_pixel_task':False}
-    #checkpoint = torch.load('deshabby.pt')
-    #model.load_state_dict(checkpoint['params'])
-    #model.eval()
-    model = torch.jit.load('deshabby.pt') 
+
+    parameters = {'inp_channels':3, 'out_channels':3, 'dim':48, 'num_blocks':[4,6,6,8], 'num_refinement_blocks':4, 'heads':[1,2,4,8], 'ffn_expansion_factor':2.66, 'bias':False, 'LayerNorm_type':'WithBias', 'dual_pixel_task':False}
+    load_arch = run_path('/content/drive/MyDrive/colab/denoising_dirty_documents/Restormer/basicsr/models/archs/restormer_arch.py')
+    model = load_arch['Restormer'](**parameters)
     
+    checkpoint = torch.load('deshabby.pt')
+    model.load_state_dict(checkpoint['params'])
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
     model = model.to(device)
     model.eval()