fixes for latest pytorch

cutie/model/big_modules.py

@@ -284,7 +284,7 @@ def forward(self,
 
             p8 = self.up_16_8(p16, f8)
             p4 = self.up_8_4(p8, f4)
-            with torch.cuda.amp.autocast(enabled=False):
+            with torch.amp.autocast('cuda',enabled=False):
                 logits = self.pred(F.relu(p4.flatten(start_dim=0, end_dim=1).float()))
 
             if update_sensory:

cutie/model/modules.py

@@ -59,7 +59,7 @@ def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
         g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
             self.g4_conv(downsample_groups(g[2], ratio=1/4))
 
-        with torch.cuda.amp.autocast(enabled=False):
+        with torch.amp.autocast('cuda', enabled=False):
             g = g.float()
             h = h.float()
             values = self.transform(torch.cat([g, h], dim=2))
@@ -76,7 +76,7 @@ def __init__(self, f_dim: int, sensory_dim: int):
         nn.init.xavier_normal_(self.transform.weight)
 
     def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
-        with torch.cuda.amp.autocast(enabled=False):
+        with torch.amp.autocast('cuda', enabled=False):
             g = g.float()
             h = h.float()
             values = self.transform(torch.cat([g, h], dim=2))

cutie/model/transformer/object_summarizer.py

@@ -75,7 +75,7 @@ def forward(self,
             pe = self.pos_enc(value)
             value = value + pe
 
-        with torch.cuda.amp.autocast(enabled=False):
+        with torch.amp.autocast('cuda', enabled=False):
             value = value.float()
             feature = self.feature_pred(value)
             logits = self.weights_pred(value)

cutie/utils/tensor_utils.py

@@ -45,7 +45,7 @@ def unpad(img: torch.Tensor, pad: Iterable[int]) -> torch.Tensor:
 
 # @torch.jit.script
 def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
-    with torch.cuda.amp.autocast(enabled=False):
+    with torch.amp.autocast('cuda',enabled=False):
         prob = prob.float()
         new_prob = torch.cat([torch.prod(1 - prob, dim=dim, keepdim=True), prob],
                              dim).clamp(1e-7, 1 - 1e-7)

gui/main_controller.py

@@ -137,7 +137,7 @@ def __init__(self, cfg: DictConfig) -> None:
     def initialize_networks(self) -> None:
         download_models_if_needed()
         self.cutie = CUTIE(self.cfg).eval().to(self.device)
-        model_weights = torch.load(self.cfg.weights, map_location=self.device)
+        model_weights = torch.load(self.cfg.weights, map_location=self.device, weights_only=True)
         self.cutie.load_weights(model_weights)
         if self.cfg.ritm_use_anime is True:
             self.click_ctrl = ClickController(self.cfg.ritm_anime_weights, self.cfg.ritm_max_size, self.cfg.ritm_zoom_size, self.cfg.ritm_expansion_ratio, device=self.device)

gui/ritm/inference/utils.py

@@ -19,7 +19,7 @@ def get_time_metrics(all_ious, elapsed_time):
 
 def load_is_model(checkpoint, device, **kwargs):
     if isinstance(checkpoint, (str, Path)):
-        state_dict = torch.load(checkpoint, map_location='cpu')
+        state_dict = torch.load(checkpoint, map_location='cpu', weights_only=True)
     else:
         state_dict = checkpoint
 

install_dependencies.bat

@@ -13,13 +13,13 @@ if errorlevel 3 (
     echo Uninstalling existing Pytorch
     .\python-3.11.7-embed-amd64\python.exe -m pip uninstall torch torchvision
     echo Installing CUDA version of PyTorch
-    .\python-3.11.7-embed-amd64\python.exe -m pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118
+    .\python-3.11.7-embed-amd64\python.exe -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu118
 	.\python-3.11.7-embed-amd64\python.exe -m pip install -r .\requirements.txt
 ) else (
     echo Uninstalling existing Pytorch
     .\python-3.11.7-embed-amd64\python.exe -m pip uninstall torch torchvision
     echo Installing CPU version of PyTorch
-    .\python-3.11.7-embed-amd64\python.exe -m pip install torch==2.3.1 torchvision==0.18.1
+    .\python-3.11.7-embed-amd64\python.exe -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
 	.\python-3.11.7-embed-amd64\python.exe -m pip install -r .\requirements.txt
 )