init

app.py

@@ -52,7 +52,7 @@ def load_checkpoint(model, ckpt_path):
         new_state_dict[k] = v
 
     model.load_state_dict(new_state_dict)
-    model.cuda()
+    # model.cuda()
     del checkpoint
     torch.cuda.empty_cache()
 
@@ -77,8 +77,8 @@ def image_infer(img_PIL):
         trainer.model.eval()
         img_lr, img_lr_up = get_img_data(img_PIL, hparams, sr_scale=4)
 
-        img_lr = img_lr.to('cuda')
-        img_lr_up = img_lr_up.to('cuda')
+        # img_lr = img_lr.to('cuda')
+        # img_lr_up = img_lr_up.to('cuda')
 
         img_sr, _ = trainer.model.sample(img_lr, img_lr_up, img_lr_up.shape)
 
@@ -89,25 +89,21 @@ def image_infer(img_PIL):
     return img_sr
 
 
-# cheetah = os.path.join(os.path.dirname(__file__), "images/cheetah1.jpg")
-
 root_path = os.path.dirname(__file__)
 
-cheetah = os.path.join(root_path, "images/lion.jpg")
+cheetah = os.path.join(root_path, "images/0801x4.png")
 print(cheetah)
 
+ckpt_path = os.path.join(root_path, 'sam_diffsr/weight/model_ckpt_steps_400000.ckpt')
+trainer = model_init(ckpt_path)
 demo = gr.Interface(image_infer, gr.Image(type="pil", value=cheetah), "image",
                     # flagging_options=["blurry", "incorrect", "other"],
                     examples=[
                         os.path.join(root_path, "images/0801x4.png"),
-                        os.path.join(root_path, "images/0809x4.png"),
+                        os.path.join(root_path, "images/0804x4.png"),
                         os.path.join(root_path, "images/0809x4.png"),
                     ]
                     )
 
 if __name__ == "__main__":
-    parent_path = Path(__file__).absolute().parent
-    fill_root = os.path.abspath(parent_path)
-    ckpt_path = os.path.join(fill_root, 'sam_diffsr/weight/model_ckpt_steps_400000.ckpt')
-    trainer = model_init(ckpt_path)
     demo.launch()