From a7ef3e04ea0b23e7967086aa6377dcfa77610073 Mon Sep 17 00:00:00 2001
From: doctorpangloss <2229300+doctorpangloss@users.noreply.github.com>
Date: Tue, 9 Dec 2025 09:05:18 -0800
Subject: [PATCH] update extra nodes

---
 comfy/model_patcher.py                 | 2 +-
 comfy_extras/nodes/nodes_dataset.py    | 4 ++--
 comfy_extras/nodes/nodes_kandinsky5.py | 4 ++--
 comfy_extras/nodes/nodes_rope.py       | 3 ++-
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index f79bbe880..2433c076e 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -340,7 +340,7 @@ class ModelPatcher(ModelManageable, PatchSupport):
     def loaded_size(self):
         return self._memory_measurements.model_loaded_weight_memory
 
-    def clone(self):
+    def clone(self) -> "ModelPatcher":
         n = self.__class__(self.model, self.load_device, self.offload_device, self.model_size(), weight_inplace_update=self.weight_inplace_update)
         n._memory_measurements = self._memory_measurements
         n.ckpt_name = self.ckpt_name
diff --git a/comfy_extras/nodes/nodes_dataset.py b/comfy_extras/nodes/nodes_dataset.py
index 4789d7d53..c184b4a5b 100644
--- a/comfy_extras/nodes/nodes_dataset.py
+++ b/comfy_extras/nodes/nodes_dataset.py
@@ -7,8 +7,8 @@ import torch
 from PIL import Image
 from typing_extensions import override
 
-import folder_paths
-import node_helpers
+from comfy.cmd import folder_paths
+from comfy import node_helpers
 from comfy_api.latest import ComfyExtension, io
 
 
diff --git a/comfy_extras/nodes/nodes_kandinsky5.py b/comfy_extras/nodes/nodes_kandinsky5.py
index 9cb234be1..4679baee7 100644
--- a/comfy_extras/nodes/nodes_kandinsky5.py
+++ b/comfy_extras/nodes/nodes_kandinsky5.py
@@ -1,5 +1,5 @@
-import nodes
-import node_helpers
+from comfy.nodes import base_nodes as nodes
+from comfy import node_helpers
 import torch
 import comfy.model_management
 import comfy.utils
diff --git a/comfy_extras/nodes/nodes_rope.py b/comfy_extras/nodes/nodes_rope.py
index d1feb031e..03faf19b1 100644
--- a/comfy_extras/nodes/nodes_rope.py
+++ b/comfy_extras/nodes/nodes_rope.py
@@ -1,3 +1,4 @@
+from comfy.model_patcher import ModelPatcher
 from comfy_api.latest import ComfyExtension, io
 from typing_extensions import override
 
@@ -29,7 +30,7 @@ class ScaleROPE(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model, scale_x, shift_x, scale_y, shift_y, scale_t, shift_t) -> io.NodeOutput:
+    def execute(cls, model: ModelPatcher, scale_x, shift_x, scale_y, shift_y, scale_t, shift_t) -> io.NodeOutput:
         m = model.clone()
         m.set_model_rope_options(scale_x, shift_x, scale_y, shift_y, scale_t, shift_t)
         return io.NodeOutput(m)