From 76c2f4f1351377dc0c3f4ee795887303600dfd04 Mon Sep 17 00:00:00 2001
From: Peter Schroedl <peter_schroedl@me.com>
Date: Wed, 27 Nov 2024 21:30:37 +0100
Subject: [PATCH] sam2 -> sam2_realtime to disambiguate package imports

---
 .gitignore                                     |  1 +
 nodes.py                                       | 18 +++++++++++++-----
 sam2/__init__.py                               |  9 ---------
 sam2_realtime/__init__.py                      |  5 +++++
 {sam2 => sam2_realtime}/build_sam.py           |  0
 .../csrc/connected_components.cu               |  0
 {sam2 => sam2_realtime}/modeling/__init__.py   |  0
 .../modeling/backbones/__init__.py             |  0
 .../modeling/backbones/hieradet.py             |  4 ++--
 .../modeling/backbones/image_encoder.py        |  0
 .../modeling/backbones/utils.py                |  0
 .../modeling/memory_attention.py               |  4 ++--
 .../modeling/memory_encoder.py                 |  2 +-
 .../modeling/position_encoding.py              |  0
 .../modeling/sam/__init__.py                   |  0
 .../modeling/sam/mask_decoder.py               |  2 +-
 .../modeling/sam/prompt_encoder.py             |  4 ++--
 .../modeling/sam/transformer.py                |  6 +++---
 {sam2 => sam2_realtime}/modeling/sam2_base.py  |  8 ++++----
 {sam2 => sam2_realtime}/modeling/sam2_utils.py |  0
 .../sam2_tensor_predictor.py                   | 15 ++++-----------
 {sam2 => sam2_realtime}/utils/__init__.py      |  0
 {sam2 => sam2_realtime}/utils/amg.py           |  0
 {sam2 => sam2_realtime}/utils/misc.py          |  0
 {sam2 => sam2_realtime}/utils/transforms.py    |  0
 25 files changed, 38 insertions(+), 40 deletions(-)
 create mode 100644 .gitignore
 delete mode 100644 sam2/__init__.py
 create mode 100644 sam2_realtime/__init__.py
 rename {sam2 => sam2_realtime}/build_sam.py (100%)
 rename {sam2 => sam2_realtime}/csrc/connected_components.cu (100%)
 rename {sam2 => sam2_realtime}/modeling/__init__.py (100%)
 rename {sam2 => sam2_realtime}/modeling/backbones/__init__.py (100%)
 rename {sam2 => sam2_realtime}/modeling/backbones/hieradet.py (98%)
 rename {sam2 => sam2_realtime}/modeling/backbones/image_encoder.py (100%)
 rename {sam2 => sam2_realtime}/modeling/backbones/utils.py (100%)
 rename {sam2 => sam2_realtime}/modeling/memory_attention.py (97%)
 rename {sam2 => sam2_realtime}/modeling/memory_encoder.py (98%)
 rename {sam2 => sam2_realtime}/modeling/position_encoding.py (100%)
 rename {sam2 => sam2_realtime}/modeling/sam/__init__.py (100%)
 rename {sam2 => sam2_realtime}/modeling/sam/mask_decoder.py (99%)
 rename {sam2 => sam2_realtime}/modeling/sam/prompt_encoder.py (98%)
 rename {sam2 => sam2_realtime}/modeling/sam/transformer.py (98%)
 rename {sam2 => sam2_realtime}/modeling/sam2_base.py (99%)
 rename {sam2 => sam2_realtime}/modeling/sam2_utils.py (100%)
 rename sam2/sam2_camera_predictor.py => sam2_realtime/sam2_tensor_predictor.py (99%)
 rename {sam2 => sam2_realtime}/utils/__init__.py (100%)
 rename {sam2 => sam2_realtime}/utils/amg.py (100%)
 rename {sam2 => sam2_realtime}/utils/misc.py (100%)
 rename {sam2 => sam2_realtime}/utils/transforms.py (100%)

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..ba0430d
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__/
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
index bc6cb0d..b6b6057 100644
--- a/nodes.py
+++ b/nodes.py
@@ -5,7 +5,15 @@ import numpy as np
 import logging
 import json
 import ast
-from .sam2.sam2_camera_predictor import SAM2CameraPredictor
+
+import sys
+
+# Add the directory containing 'sam2_realtime' to sys.path
+current_directory = os.path.dirname(os.path.abspath(__file__))
+sam2_realtime_path = os.path.join(current_directory)  # Adjust the relative path
+sys.path.append(sam2_realtime_path)
+print("sys.path:", sys.path)
+from sam2_realtime.sam2_tensor_predictor import SAM2TensorPredictor
 from comfy.utils import load_torch_file
 
 from omegaconf import OmegaConf
@@ -40,7 +48,7 @@ class DownloadAndLoadSAM2RealtimeModel:
     RETURN_TYPES = ("SAM2MODEL",)
     RETURN_NAMES = ("sam2_model",)
     FUNCTION = "loadmodel"
-    CATEGORY = "SAM2-Realtime"
+    CATEGORY = "sam2_realtime"
 
     def loadmodel(self, model, segmentor, device, precision):
         if precision != 'fp32' and device == 'cpu':
@@ -80,7 +88,7 @@ class DownloadAndLoadSAM2RealtimeModel:
             cfg = compose(config_name=model_cfg)
 
             hydra_overrides = [
-                "++model._target_=sam2.sam2_camera_predictor.SAM2CameraPredictor",
+                "++model._target_=sam2_realtime.sam2_tensor_predictor.SAM2TensorPredictor",
             ]
             hydra_overrides_extra = [
                 "++model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability=true",
@@ -146,7 +154,7 @@ class Sam2RealtimeSegmentation:
     RETURN_NAMES = ("PROCESSED_IMAGES","MASK",)
     RETURN_TYPES = ("IMAGE", "IMAGE",)
     FUNCTION = "segment_images"
-    CATEGORY = "SAM2-Realtime"
+    CATEGORY = "sam2_realtime"
 
     def __init__(self):
         self.predictor = None
@@ -240,6 +248,6 @@ NODE_CLASS_MAPPINGS = {
     "Sam2RealtimeSegmentation": Sam2RealtimeSegmentation
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "DownloadAndLoadSAM2RealtimeModel": "(Down)Load SAM2-Realtime Model",
+    "DownloadAndLoadSAM2RealtimeModel": "(Down)Load sam2_realtime Model",
     "Sam2RealtimeSegmentation": "Sam2RealtimeSegmentation"
 }
diff --git a/sam2/__init__.py b/sam2/__init__.py
deleted file mode 100644
index 4f1c77e..0000000
--- a/sam2/__init__.py
+++ /dev/null
@@ -1,9 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-# from hydra import initialize_config_module
-
-# initialize_config_module("sam2_configs", version_base="1.2")
diff --git a/sam2_realtime/__init__.py b/sam2_realtime/__init__.py
new file mode 100644
index 0000000..21db546
--- /dev/null
+++ b/sam2_realtime/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
\ No newline at end of file
diff --git a/sam2/build_sam.py b/sam2_realtime/build_sam.py
similarity index 100%
rename from sam2/build_sam.py
rename to sam2_realtime/build_sam.py
diff --git a/sam2/csrc/connected_components.cu b/sam2_realtime/csrc/connected_components.cu
similarity index 100%
rename from sam2/csrc/connected_components.cu
rename to sam2_realtime/csrc/connected_components.cu
diff --git a/sam2/modeling/__init__.py b/sam2_realtime/modeling/__init__.py
similarity index 100%
rename from sam2/modeling/__init__.py
rename to sam2_realtime/modeling/__init__.py
diff --git a/sam2/modeling/backbones/__init__.py b/sam2_realtime/modeling/backbones/__init__.py
similarity index 100%
rename from sam2/modeling/backbones/__init__.py
rename to sam2_realtime/modeling/backbones/__init__.py
diff --git a/sam2/modeling/backbones/hieradet.py b/sam2_realtime/modeling/backbones/hieradet.py
similarity index 98%
rename from sam2/modeling/backbones/hieradet.py
rename to sam2_realtime/modeling/backbones/hieradet.py
index 1ae7d4c..d651a08 100644
--- a/sam2/modeling/backbones/hieradet.py
+++ b/sam2_realtime/modeling/backbones/hieradet.py
@@ -11,13 +11,13 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from sam2.modeling.backbones.utils import (
+from sam2_realtime.modeling.backbones.utils import (
     PatchEmbed,
     window_partition,
     window_unpartition,
 )
 
-from sam2.modeling.sam2_utils import DropPath, MLP
+from sam2_realtime.modeling.sam2_utils import DropPath, MLP
 
 
 def do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) -> torch.Tensor:
diff --git a/sam2/modeling/backbones/image_encoder.py b/sam2_realtime/modeling/backbones/image_encoder.py
similarity index 100%
rename from sam2/modeling/backbones/image_encoder.py
rename to sam2_realtime/modeling/backbones/image_encoder.py
diff --git a/sam2/modeling/backbones/utils.py b/sam2_realtime/modeling/backbones/utils.py
similarity index 100%
rename from sam2/modeling/backbones/utils.py
rename to sam2_realtime/modeling/backbones/utils.py
diff --git a/sam2/modeling/memory_attention.py b/sam2_realtime/modeling/memory_attention.py
similarity index 97%
rename from sam2/modeling/memory_attention.py
rename to sam2_realtime/modeling/memory_attention.py
index 0b07f9d..13df4bc 100644
--- a/sam2/modeling/memory_attention.py
+++ b/sam2_realtime/modeling/memory_attention.py
@@ -9,9 +9,9 @@ from typing import Optional
 import torch
 from torch import nn, Tensor
 
-from sam2.modeling.sam.transformer import RoPEAttention
+from sam2_realtime.modeling.sam.transformer import RoPEAttention
 
-from sam2.modeling.sam2_utils import get_activation_fn, get_clones
+from sam2_realtime.modeling.sam2_utils import get_activation_fn, get_clones
 
 
 class MemoryAttentionLayer(nn.Module):
diff --git a/sam2/modeling/memory_encoder.py b/sam2_realtime/modeling/memory_encoder.py
similarity index 98%
rename from sam2/modeling/memory_encoder.py
rename to sam2_realtime/modeling/memory_encoder.py
index f60202d..8848924 100644
--- a/sam2/modeling/memory_encoder.py
+++ b/sam2_realtime/modeling/memory_encoder.py
@@ -11,7 +11,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from sam2.modeling.sam2_utils import DropPath, get_clones, LayerNorm2d
+from sam2_realtime.modeling.sam2_utils import DropPath, get_clones, LayerNorm2d
 
 
 class MaskDownSampler(nn.Module):
diff --git a/sam2/modeling/position_encoding.py b/sam2_realtime/modeling/position_encoding.py
similarity index 100%
rename from sam2/modeling/position_encoding.py
rename to sam2_realtime/modeling/position_encoding.py
diff --git a/sam2/modeling/sam/__init__.py b/sam2_realtime/modeling/sam/__init__.py
similarity index 100%
rename from sam2/modeling/sam/__init__.py
rename to sam2_realtime/modeling/sam/__init__.py
diff --git a/sam2/modeling/sam/mask_decoder.py b/sam2_realtime/modeling/sam/mask_decoder.py
similarity index 99%
rename from sam2/modeling/sam/mask_decoder.py
rename to sam2_realtime/modeling/sam/mask_decoder.py
index b7c7dfd..f31a941 100644
--- a/sam2/modeling/sam/mask_decoder.py
+++ b/sam2_realtime/modeling/sam/mask_decoder.py
@@ -9,7 +9,7 @@ from typing import List, Optional, Tuple, Type
 import torch
 from torch import nn
 
-from sam2.modeling.sam2_utils import LayerNorm2d, MLP
+from sam2_realtime.modeling.sam2_utils import LayerNorm2d, MLP
 
 
 class MaskDecoder(nn.Module):
diff --git a/sam2/modeling/sam/prompt_encoder.py b/sam2_realtime/modeling/sam/prompt_encoder.py
similarity index 98%
rename from sam2/modeling/sam/prompt_encoder.py
rename to sam2_realtime/modeling/sam/prompt_encoder.py
index 6b3bbb9..7157a7a 100644
--- a/sam2/modeling/sam/prompt_encoder.py
+++ b/sam2_realtime/modeling/sam/prompt_encoder.py
@@ -9,9 +9,9 @@ from typing import Optional, Tuple, Type
 import torch
 from torch import nn
 
-from sam2.modeling.position_encoding import PositionEmbeddingRandom
+from sam2_realtime.modeling.position_encoding import PositionEmbeddingRandom
 
-from sam2.modeling.sam2_utils import LayerNorm2d
+from sam2_realtime.modeling.sam2_utils import LayerNorm2d
 
 
 class PromptEncoder(nn.Module):
diff --git a/sam2/modeling/sam/transformer.py b/sam2_realtime/modeling/sam/transformer.py
similarity index 98%
rename from sam2/modeling/sam/transformer.py
rename to sam2_realtime/modeling/sam/transformer.py
index 2a0b796..46faee7 100644
--- a/sam2/modeling/sam/transformer.py
+++ b/sam2_realtime/modeling/sam/transformer.py
@@ -13,10 +13,10 @@ import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
 
-from sam2.modeling.position_encoding import apply_rotary_enc, compute_axial_cis
+from sam2_realtime.modeling.position_encoding import apply_rotary_enc, compute_axial_cis
 
-from sam2.modeling.sam2_utils import MLP
-from sam2.utils.misc import get_sdpa_settings
+from sam2_realtime.modeling.sam2_utils import MLP
+from sam2_realtime.utils.misc import get_sdpa_settings
 
 warnings.simplefilter(action="ignore", category=FutureWarning)
 OLD_GPU, USE_FLASH_ATTN, MATH_KERNEL_ON = get_sdpa_settings()
diff --git a/sam2/modeling/sam2_base.py b/sam2_realtime/modeling/sam2_base.py
similarity index 99%
rename from sam2/modeling/sam2_base.py
rename to sam2_realtime/modeling/sam2_base.py
index 2b5251f..700591f 100644
--- a/sam2/modeling/sam2_base.py
+++ b/sam2_realtime/modeling/sam2_base.py
@@ -10,10 +10,10 @@ import torch.nn.functional as F
 
 from torch.nn.init import trunc_normal_
 
-from sam2.modeling.sam.mask_decoder import MaskDecoder
-from sam2.modeling.sam.prompt_encoder import PromptEncoder
-from sam2.modeling.sam.transformer import TwoWayTransformer
-from sam2.modeling.sam2_utils import get_1d_sine_pe, MLP, select_closest_cond_frames
+from sam2_realtime.modeling.sam.mask_decoder import MaskDecoder
+from sam2_realtime.modeling.sam.prompt_encoder import PromptEncoder
+from sam2_realtime.modeling.sam.transformer import TwoWayTransformer
+from sam2_realtime.modeling.sam2_utils import get_1d_sine_pe, MLP, select_closest_cond_frames
 
 # a large negative value as a placeholder score for missing objects
 NO_OBJ_SCORE = -1024.0
diff --git a/sam2/modeling/sam2_utils.py b/sam2_realtime/modeling/sam2_utils.py
similarity index 100%
rename from sam2/modeling/sam2_utils.py
rename to sam2_realtime/modeling/sam2_utils.py
diff --git a/sam2/sam2_camera_predictor.py b/sam2_realtime/sam2_tensor_predictor.py
similarity index 99%
rename from sam2/sam2_camera_predictor.py
rename to sam2_realtime/sam2_tensor_predictor.py
index ff9b40b..589210d 100644
--- a/sam2/sam2_camera_predictor.py
+++ b/sam2_realtime/sam2_tensor_predictor.py
@@ -7,22 +7,15 @@
 from collections import OrderedDict
 
 import torch
+import numpy as np
 
 from tqdm import tqdm
 
-import sys
-import os
-
-# To import from local sam2
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-
-from sam2.modeling.sam2_base import NO_OBJ_SCORE, SAM2Base
-from sam2.utils.misc import concat_points, fill_holes_in_mask_scores, load_video_frames
-import numpy as np
-import cv2
+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 SAM2CameraPredictor(SAM2Base):
+class SAM2TensorPredictor(SAM2Base):
     """The predictor class to handle user interactions and manage inference states."""
 
     def __init__(
diff --git a/sam2/utils/__init__.py b/sam2_realtime/utils/__init__.py
similarity index 100%
rename from sam2/utils/__init__.py
rename to sam2_realtime/utils/__init__.py
diff --git a/sam2/utils/amg.py b/sam2_realtime/utils/amg.py
similarity index 100%
rename from sam2/utils/amg.py
rename to sam2_realtime/utils/amg.py
diff --git a/sam2/utils/misc.py b/sam2_realtime/utils/misc.py
similarity index 100%
rename from sam2/utils/misc.py
rename to sam2_realtime/utils/misc.py
diff --git a/sam2/utils/transforms.py b/sam2_realtime/utils/transforms.py
similarity index 100%
rename from sam2/utils/transforms.py
rename to sam2_realtime/utils/transforms.py