remove import not needed

Relax python 3.11.15 -> 3.10.15

requirements.txt

@@ -4,4 +4,4 @@ tqdm>=4.66.1
 hydra-core>=1.3.2
 iopath>=0.1.10
 pillow>=9.4.0
--e .
+git+https://github.com/pschroedl/ComfyUI-SAM2-Realtime.git@main#egg=sam2_realtime

sam2_realtime/sam2_tensor_predictor.py

@@ -14,7 +14,6 @@ from tqdm import tqdm
 from sam2_realtime.modeling.sam2_base import NO_OBJ_SCORE, SAM2Base
 from sam2_realtime.utils.misc import concat_points, fill_holes_in_mask_scores, load_video_frames
 
-
 class SAM2TensorPredictor(SAM2Base):
     """The predictor class to handle user interactions and manage inference states."""
 
@@ -55,7 +54,9 @@ class SAM2TensorPredictor(SAM2Base):
             img = img.float()
         else:
             raise ValueError("Input must be a numpy array or a PyTorch tensor")
-
+        #save original height/width
+        orig_h, orig_w = img.shape[1:]
+        
         # Resize to the target size (supports tensor resizing)
         img = torch.nn.functional.interpolate(
             img.unsqueeze(0), size=(image_size, image_size), mode="bilinear", align_corners=False
@@ -68,28 +69,29 @@ class SAM2TensorPredictor(SAM2Base):
         img /= img_std
 
         height, width = img.shape[1:]  # CHW format
-        return img, width, height
+        return img, width, height, orig_w, orig_h
 
     @torch.inference_mode()
     def load_first_frame(self, img):
         if isinstance(img, torch.Tensor):
             img = img.to(self.device)  # Ensure the tensor is on the correct device
-
+            
         self.condition_state = self._init_state(
             offload_video_to_cpu=False, offload_state_to_cpu=False
         )
-        img, width, height = self.prepare_data(img, image_size=self.image_size)
+        img, width, height, orig_w, orig_h = self.prepare_data(img, image_size=self.image_size)
+        self._orig_hw = (orig_w, orig_h)
         self.condition_state["images"] = [img]
         self.condition_state["num_frames"] = len(self.condition_state["images"])
         self.condition_state["video_height"] = height
         self.condition_state["video_width"] = width
         self._get_image_feature(frame_idx=0, batch_size=1)
-
+    
     def add_conditioning_frame(self, img):
         if isinstance(img, torch.Tensor):
             img = img.to(self.device)  # Ensure the tensor is on the correct device
 
-        img, width, height = self.prepare_data(img, image_size=self.image_size)
+        img, width, height, _, _ = self.prepare_data(img, image_size=self.image_size)
         self.condition_state["images"].append(img)
         self.condition_state["num_frames"] = len(self.condition_state["images"])
         self._get_image_feature(
@@ -235,14 +237,15 @@ class SAM2TensorPredictor(SAM2Base):
                     points = torch.cat([box_coords, points], dim=1)
                     labels = torch.cat([box_labels, labels], dim=1)
         if normalize_coords:
-            video_H = self.condition_state["video_height"]
-            video_W = self.condition_state["video_width"]
-            points = points / torch.tensor([video_W, video_H]).to(points.device)
+            #video_H = self.condition_state["video_height"]
+            #video_W = self.condition_state["video_width"]
+            orig_w, orig_h = self._orig_hw
+            
+            points = points / torch.tensor([orig_w, orig_h]).to(points.device)
         # scale the (normalized) coordinates by the model's internal image size
         points = points * self.image_size
         points = points.to(self.condition_state["device"])
         labels = labels.to(self.condition_state["device"])
-
         if not clear_old_points:
             point_inputs = point_inputs_per_frame.get(frame_idx, None)
         else:
@@ -342,14 +345,16 @@ class SAM2TensorPredictor(SAM2Base):
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)  # add batch dimension
         if normalize_coords:
-            video_H = self.condition_state["video_height"]
-            video_W = self.condition_state["video_width"]
-            points = points / torch.tensor([video_W, video_H]).to(points.device)
+            #video_H = self.condition_state["video_height"]
+            #video_W = self.condition_state["video_width"]
+            orig_w, orig_h = self._orig_hw
+            
+            points = points / torch.tensor([orig_w, orig_h]).to(points.device)
+            
         # scale the (normalized) coordinates by the model's internal image size
         points = points * self.image_size
         points = points.to(self.condition_state["device"])
         labels = labels.to(self.condition_state["device"])
-
         if not clear_old_points:
             point_inputs = point_inputs_per_frame.get(frame_idx, None)
         else:
@@ -769,7 +774,7 @@ class SAM2TensorPredictor(SAM2Base):
         if isinstance(img, torch.Tensor):
             img = img.to(self.device)  # Ensure the tensor is on the correct device
 
-        img, _, _ = self.prepare_data(img, image_size=self.image_size)
+        img, _, _ , _, _ = self.prepare_data(img, image_size=self.image_size)
 
         output_dict = self.condition_state["output_dict"]
         obj_ids = self.condition_state["obj_ids"]

setup.py

@@ -4,6 +4,10 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+# Set the CUDA architecture list
+os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0 8.6+PTX 8.7 9.0 9.0a"
+
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -53,7 +57,7 @@ setup(
     license=LICENSE,
     packages=find_packages(),
     install_requires=REQUIRED_PACKAGES,
-    python_requires=">=3.11.10",
+    python_requires=">=3.10.15",
     ext_modules=get_extensions(),
     cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
 )