add compile call

gradio_app.py

@@ -6,6 +6,7 @@ import torch.nn.functional as F
 import os
 import time
 import spaces
+import subprocess
 
 from dataloader.stereo import transforms
 from utils.utils import InputPadder, calc_noc_mask
@@ -37,11 +38,18 @@ class MatchStereoDemo:
         
     def load_model(self, mode, variant, precision, mat_impl):
         """load model, skip if the model has been loaded"""
+        current_has_cuda = torch.cuda.is_available()
+        if current_has_cuda != self.has_cuda:
+            print(f"CUDA status changed: {self.has_cuda} -> {current_has_cuda}")
+            self.has_cuda = current_has_cuda
+            self.device = "cuda" if self.has_cuda else 'cpu'
+            
         if (self.model is not None and 
             self.current_variant == variant and 
             self.current_mode == mode and
             self.current_precision == precision and
-            self.current_mat_impl == mat_impl):
+            self.current_mat_impl == mat_impl and
+            self.has_cuda == current_has_cuda):
             return "Model already loaded"
             
         # fixed checkpoint path
@@ -65,6 +73,7 @@ class MatchStereoDemo:
         
         if not self.has_cuda:
             precision = "fp32"
+            mat_impl = "pytorch"
         dtypes = {'fp32': torch.float32, 'fp16': torch.float16}
         self.dtype = dtypes[precision]
         
@@ -138,6 +147,12 @@ class MatchStereoDemo:
     def process_images(self, left_image, right_image, mode, variant, 
                       low_res_init=False, inference_size_name="Original", 
                       precision="fp32", mat_impl="pytorch"):
+        current_has_cuda = torch.cuda.is_available()
+        if current_has_cuda != self.has_cuda:
+            print(f"CUDA status changed before processing: {self.has_cuda} -> {current_has_cuda}")
+            self.has_cuda = current_has_cuda
+            self.device = "cuda" if self.has_cuda else 'cpu'
+            
         if not self.has_cuda:
             precision = "fp32"
             mat_impl = "pytorch"
@@ -247,6 +262,37 @@ class MatchStereoDemo:
 
 demo_model = MatchStereoDemo()
 
+def compile_cuda_extensions():
+    try:
+        print("Start compiling CUDA extension...")
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        models_dir = os.path.join(current_dir, "models")
+        compile_script = os.path.join(models_dir, "compile.sh")
+        
+        if os.path.exists(compile_script):
+            original_cwd = os.getcwd()
+            os.chdir(models_dir)
+            
+            result = subprocess.run(["bash", "compile.sh"], 
+                                  capture_output=True, text=True)
+            
+            os.chdir(original_cwd)
+            
+            if result.returncode == 0:
+                print("CUDA extension compile succeed!")
+                print("output:", result.stdout)
+            else:
+                print("CUDA extension compile failed!")
+                print(result.stderr)
+                print(result.stdout)
+        else:
+            print(f"no compile scripts found: {compile_script}")
+            
+    except Exception as e:
+        print(f"Error during compile: {e}")
+
+compile_cuda_extensions()
+
 # example images
 examples = [
     ["examples/booster_bathroom_left.png", "examples/booster_bathroom_right.png", "stereo", "tiny"],
@@ -280,8 +326,11 @@ with gr.Blocks(title="MatchStereo/MatchFlow Demo") as demo:
     gr.Markdown("# MatchStereo/MatchFlow Demo")
     gr.Markdown("Upload stereo images for disparity estimation or consecutive frames for optical flow estimation.")
     
-    if not demo_model.has_cuda:
+    current_has_cuda = torch.cuda.is_available()
+    if not current_has_cuda:
         gr.Markdown("> Note: Running on CPU. Some options (fp16, cuda) are disabled.")
+    else:
+        gr.Markdown(f"> Note: Running on GPU ({torch.cuda.get_device_name(0)}).")
     
     with gr.Row():
         with gr.Column():
@@ -321,14 +370,14 @@ with gr.Blocks(title="MatchStereo/MatchFlow Demo") as demo:
                     label="Precision",
                     value="fp32",
                     info="Model precision",
-                    interactive=demo_model.has_cuda
+                    interactive=current_has_cuda
                 )
                 mat_impl = gr.Radio(
                     choices=["cuda", "pytorch"],
                     label="MatchAttention Implementation",
                     value="cuda",
                     info="MatchAttention implementations",
-                    interactive=demo_model.has_cuda
+                    interactive=current_has_cuda
                 )
             
             run_btn = gr.Button("Run Inference", variant="primary")