fixed checkpoint loading requires GPU issue

app.py

@@ -19,7 +19,7 @@ if RUN_MODE != "local":
 ## step 1: set up model
 device = "cpu"
 checkpt_path = "checkpoints/model_best.pth.tar"
-invhalfer = Inferencer(checkpoint_path=checkpt_path, model=ResHalf(train=False), use_cuda=False)
+invhalfer = Inferencer(checkpoint_path=checkpt_path, model=ResHalf(train=False), use_cuda=False, multi_gpu=False)
 
 
 def prepare_data(input_img, decoding_only=False):
@@ -41,7 +41,7 @@ def run_invhalf(invhalfer, input_img, decoding_only, device="cuda"):
         print('>>>:halftoning mode')
         resHalftone, resColor = invhalfer(input_img, decoding_only=decoding_only)
         output = util.tensor2img(resHalftone / 2. + 0.5) * 255.
-    return (output+0.5).astype(np.uint8)
+    return np.clip(output, 0, 255).astype(np.uint8)
 
 
 def click_run(input_img, decoding_only):
@@ -66,7 +66,7 @@ with demo:
 
     Button_run.click(fn=click_run, inputs=[Image_input, Radio_mode], outputs=Image_output)
     
-    if RUN_MODE == "local":
+    if RUN_MODE != "local":
         gr.Examples(examples=[
                     ['girl.png', "Halftoning (Photo2Halftone)"],
                     ['wave.png', "Halftoning (Photo2Halftone)"],
@@ -74,7 +74,7 @@ with demo:
                     ], 
                     inputs=[Image_input,Radio_mode], outputs=[Image_output], label="Examples")
 
-if RUN_MODE != "local":
+if RUN_MODE == "local":
     demo.launch(server_name='9.134.253.83',server_port=7788)
 else:
     demo.launch()

inference.py

@@ -13,18 +13,25 @@ from model.loss import l1_loss
 from utils import util
 from utils.dct import DCT_Lowfrequency
 from utils.filters_tensor import bgr2gray
-
+from collections import OrderedDict
 
 class Inferencer:
     def __init__(self, checkpoint_path, model, use_cuda=True, multi_gpu=True):
-        self.checkpoint = torch.load(checkpoint_path)
+        self.checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
         self.use_cuda = use_cuda
         self.model = model.eval()
         if multi_gpu:
             self.model = torch.nn.DataParallel(self.model)
+            state_dict = self.checkpoint['state_dict']
+        else:
+            ## remove keyword "module" in the state_dict
+            state_dict = OrderedDict()
+            for k, v in self.checkpoint['state_dict'].items():
+                name = k[7:]
+                state_dict[name] = v 
         if self.use_cuda:
             self.model = self.model.cuda()
-        self.model.load_state_dict(self.checkpoint['state_dict'])
+        self.model.load_state_dict(state_dict)
 
     def __call__(self, input_img, decoding_only=False):
         with torch.no_grad():