Update app.py

app.py

@@ -1,5 +1,3 @@
-#origin
-
 from seg import U2NETP
 from GeoTr import GeoTr
 from IllTr import IllTr
@@ -18,10 +16,6 @@ import argparse
 import warnings
 warnings.filterwarnings('ignore')
 
-
-
-
-
 import gradio as gr
 
 example_img_list = ['51_1 copy.png','48_2 copy.png','25.jpg'] 
@@ -41,48 +35,21 @@ class GeoTr_Seg(nn.Module):
         bm = (2 * (bm / 286.8) - 1) * 0.99
         
         return bm
-        
-
-def reload_model(model, path=""):
-    if not bool(path):
-        return model
-    else:
-        model_dict = model.state_dict()
-        pretrained_dict = torch.load(path, map_location='cpu')
-        #print(len(pretrained_dict.keys()))
-        pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
-        #print(len(pretrained_dict.keys()))
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
 
-        return model
-        
-
-def reload_segmodel(model, path=""):
-    if not bool(path):
-        return model
-    else:
-        model_dict = model.state_dict()
-        pretrained_dict = torch.load(path, map_location='cpu')
-        #print(len(pretrained_dict.keys()))
-        pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict}
-        #print(len(pretrained_dict.keys()))
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
-
-        return model
-        
+# Initialize models
+GeoTr_Seg_model = GeoTr_Seg()
+IllTr_model = IllTr()
 
+# Load models only once
+reload_segmodel(GeoTr_Seg_model.msk, './model_pretrained/seg.pth')
+reload_model(GeoTr_Seg_model.GeoTr, './model_pretrained/geotr.pth')
+reload_model(IllTr_model, './model_pretrained/illtr.pth')
 
+# Compile models (assuming PyTorch 2.0)
+GeoTr_Seg_model = torch.compile(GeoTr_Seg_model)
+IllTr_model = torch.compile(IllTr_model)
 
 def process_image(input_image):
-    GeoTr_Seg_model = GeoTr_Seg()#.cuda()
-    reload_segmodel(GeoTr_Seg_model.msk, './model_pretrained/seg.pth')
-    reload_model(GeoTr_Seg_model.GeoTr, './model_pretrained/geotr.pth')
-
-    IllTr_model = IllTr()#.cuda()
-    reload_model(IllTr_model, './model_pretrained/illtr.pth')
-
     GeoTr_Seg_model.eval()
     IllTr_model.eval()
 
@@ -112,13 +79,9 @@ def process_image(input_image):
         else:
             return Image.fromarray(img_geo)
 
-
-
 # Define Gradio interface
 input_image = gr.inputs.Image()
 output_image = gr.outputs.Image(type='pil')
 
-
 iface = gr.Interface(fn=process_image, inputs=input_image, outputs=output_image, title="DocTr", examples=example_img_list)
 iface.launch()
-