go back to GPU ZERO

app.py

@@ -10,22 +10,22 @@ from utils.florence import load_florence_model, run_florence_inference, \
     FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
 from utils.sam import load_sam_image_model, run_sam_inference
 
-# DEVICE = torch.device("cuda")
-DEVICE = torch.device("cpu")
+DEVICE = torch.device("cuda")
+# DEVICE = torch.device("cpu")
 
-# torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
-# if torch.cuda.get_device_properties(0).major >= 8:
-#     torch.backends.cuda.matmul.allow_tf32 = True
-#     torch.backends.cudnn.allow_tf32 = True
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
 
 
 FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE)
 SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE)
 
 
-# @spaces.GPU(duration=20)
-# @torch.inference_mode()
-# @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+@spaces.GPU(duration=20)
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process_image(image_input, text_input) -> Optional[Image.Image]:
     if not image_input:
         gr.Info("Please upload an image.")

utils/florence.py

@@ -7,8 +7,8 @@ from PIL import Image
 from transformers import AutoModelForCausalLM, AutoProcessor
 from transformers.dynamic_module_utils import get_imports
 
-FLORENCE_CHECKPOINT = "microsoft/Florence-2-base"
-# FLORENCE_CHECKPOINT = "microsoft/Florence-2-large"
+# FLORENCE_CHECKPOINT = "microsoft/Florence-2-base"
+FLORENCE_CHECKPOINT = "microsoft/Florence-2-large"
 FLORENCE_OBJECT_DETECTION_TASK = '<OD>'
 FLORENCE_DETAILED_CAPTION_TASK = '<MORE_DETAILED_CAPTION>'
 FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK = '<CAPTION_TO_PHRASE_GROUNDING>'

utils/sam.py

@@ -7,10 +7,10 @@ from PIL import Image
 from sam2.build_sam import build_sam2, build_sam2_video_predictor
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-SAM_CHECKPOINT = "checkpoints/sam2_hiera_small.pt"
-SAM_CONFIG = "sam2_hiera_s.yaml"
-# SAM_CHECKPOINT = "checkpoints/sam2_hiera_large.pt"
-# SAM_CONFIG = "sam2_hiera_l.yaml"
+# SAM_CHECKPOINT = "checkpoints/sam2_hiera_small.pt"
+# SAM_CONFIG = "sam2_hiera_s.yaml"
+SAM_CHECKPOINT = "checkpoints/sam2_hiera_large.pt"
+SAM_CONFIG = "sam2_hiera_l.yaml"
 
 
 def load_sam_image_model(