Update app.py

app.py

@@ -10,6 +10,7 @@ from net.dornet_ddp import Net_ddp
 
 # init
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print(device)
 net = Net(tiny_model=False).to(device)
 model_ckpt_map = {
     "RGB-D-D": "./checkpoints/RGBDD.pth",
@@ -17,6 +18,7 @@ model_ckpt_map = {
 }
 
 # load model
+@spaces.GPU
 def load_model(model_type: str):
     global net
     ckpt_path = model_ckpt_map[model_type]
@@ -35,6 +37,7 @@ load_model("RGB-D-D")
 
 
 # data process
+@spaces.GPU
 def preprocess_inputs(rgb_image: Image.Image, lr_depth: Image.Image):
     image = np.array(rgb_image.convert("RGB")).astype(np.float32)
     h, w, _ = image.shape
@@ -56,6 +59,7 @@ def preprocess_inputs(rgb_image: Image.Image, lr_depth: Image.Image):
 
 
 # model inference
+@spaces.GPU
 @torch.no_grad()
 def infer(rgb_image: Image.Image, lr_depth: Image.Image, model_type: str):
     load_model(model_type)  # reset weight
@@ -80,38 +84,6 @@ def infer(rgb_image: Image.Image, lr_depth: Image.Image, model_type: str):
     return pred_gray, Image.fromarray(pred_heat)
 
 
-# Gradio
-# demo = gr.Interface(
-#     fn=infer,
-#     inputs=[
-#         gr.Image(label="RGB Image", type="pil"),
-#         gr.Image(label="Low-res Depth", type="pil", image_mode="I"),
-#         gr.Dropdown(choices=["RGB-D-D", "TOFDSR"], label="Model Type", value="RGB-D-D")
-#     ],
-#     outputs=[
-#         gr.Image(label="DORNet Output", type="pil", elem_classes=["output-image"]),
-#         gr.Image(label="Normalized Output", type="pil", elem_classes=["output-image"])
-#     ],
-#     examples=[
-#         ["examples/RGB-D-D/20200518160957_RGB.jpg", "examples/RGB-D-D/20200518160957_LR_fill_depth.png", "RGB-D-D"],
-#         ["examples/TOFDSR/2020_09_08_13_59_59_435_rgb_rgb_crop.png", "examples/TOFDSR/2020_09_08_13_59_59_435_rgb_depth_crop_fill.png", "TOFDSR"],
-#     ],
-#     allow_flagging="never",
-#     title="DORNet: A Degradation Oriented and Regularized Network for Blind Depth Super-Resolution \n CVPR 2025 (Oral Presentation)",
-#     css="""
-#         .output-image {
-#             display: flex;
-#             justify-content: center;
-#             align-items: center;
-#         }
-#         .output-image img {
-#             margin: auto;
-#             display: block;
-#         }
-#         """
-# )
-#
-# demo.launch(share=True)
 Intro = """
 ## DORNet: A Degradation Oriented and Regularized Network for Blind Depth Super-Resolution
 [📄 Paper](https://arxiv.org/pdf/2410.11666) • [💻 Code](https://github.com/yanzq95/DORNet) • [📦 Model](https://huggingface.co/wzxwyx/DORNet/tree/main)