Update app.py

app.py

@@ -17,12 +17,12 @@ import gradio as gr
 import open_clip
 from sam2.build_sam import build_sam2
 from mask_adapter.modeling.meta_arch.mask_adapter_head import build_mask_adapter
+import spaces
 
 
 
 
 def setup_cfg(config_file):
-    # load config from file and command-line arguments
     cfg = get_cfg()
     add_deeplab_config(cfg)
     add_maskformer2_config(cfg)
@@ -32,7 +32,9 @@ def setup_cfg(config_file):
     cfg.freeze()
     return cfg
 
-
+@spaces.GPU
+@torch.no_grad()
+@torch.autocast(device_type="cuda", dtype=torch.float16)
 def inference_automatic(input_img, class_names):
     mp.set_start_method("spawn", force=True)
     config_file = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'
@@ -46,10 +48,11 @@ def inference_automatic(input_img, class_names):
 
     return Image.fromarray(np.uint8(visualized_output.get_image())).convert('RGB')
 
-
+@spaces.GPU
+@torch.no_grad()
+@torch.autocast(device_type="cuda", dtype=torch.float16)
 def inference_point(input_img, evt: gr.SelectData,):
-    # In point mode, implement the logic to process points from the user click (x, y)
-    # You can adjust your segmentation logic based on clicked points.
+
     x, y = evt.index[0], evt.index[1]
     points = [[x, y]]
     print(f"Selected point: {points}")
@@ -78,37 +81,32 @@ sam2_model = None
 clip_model = None
 mask_adapter = None
 
-# 加载和初始化函数
 def initialize_models(sam_path, adapter_pth, model_cfg, cfg):
     cfg = setup_cfg(cfg)
     global sam2_model, clip_model, mask_adapter
 
-    # SAM2初始化
     if sam2_model is None:
         sam2_model = build_sam2(model_cfg, sam_path, device="cuda", apply_postprocessing=False)
         print("SAM2 model initialized.")
     
-    # CLIP模型初始化
     if clip_model is None:
         clip_model, _, _ = open_clip.create_model_and_transforms("convnext_large_d_320", pretrained="laion2b_s29b_b131k_ft_soup")
+        clip_model = clip_model.eval()
         clip_model = clip_model.to("cuda")
         print("CLIP model initialized.")
     
-    # Mask Adapter模型初始化
     if mask_adapter is None:
-        mask_adapter = build_mask_adapter(cfg, "MASKAdapterHead").cuda()
-        # 加载Adapter状态字典
+        mask_adapter = build_mask_adapter(cfg, "MASKAdapterHead").to("cuda")
+        mask_adapter = mask_adapter.eval()
         adapter_state_dict = torch.load(adapter_pth)
         mask_adapter.load_state_dict(adapter_state_dict)
         print("Mask Adapter model initialized.")
 
-# 初始化配置和模型
 model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
 sam_path = './sam2.1_hiera_large.pt'
 adapter_pth = './model_0279999_with_sem_new.pth'
 cfg = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'
 
-# 调用初始化函数
 initialize_models(sam_path, adapter_pth, model_cfg, cfg)
 
 # Examples for testing