From 60abdef03ef433e5aeae4d8c3944cd075a0a7bc4 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 3 Mar 2025 16:55:56 +0200
Subject: [PATCH] rename WanVideoTeaCache -> WanVideoTeaCacheKJ to avoid clash
 with wrapper

---
 __init__.py                       |  3 ++-
 nodes/model_optimization_nodes.py |  2 +-
 nodes/nodes.py                    | 44 +++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/__init__.py b/__init__.py
index af20c95..e0d40b8 100644
--- a/__init__.py
+++ b/__init__.py
@@ -183,7 +183,8 @@ NODE_CONFIG = {
     "ScheduledCFGGuidance": {"class": ScheduledCFGGuidance, "name": "Scheduled CFG Guidance"},
     "ApplyRifleXRoPE_HunuyanVideo": {"class": ApplyRifleXRoPE_HunuyanVideo, "name": "Apply RifleXRoPE HunuyanVideo"},
     "ApplyRifleXRoPE_WanVideo": {"class": ApplyRifleXRoPE_WanVideo, "name": "Apply RifleXRoPE WanVideo"},
-    "WanVideoTeaCache": {"class": WanVideoTeaCache, "name": "WanVideo Tea Cache"},
+    "WanVideoTeaCacheKJ": {"class": WanVideoTeaCacheKJ, "name": "WanVideo Tea Cache"},
+    "TimerNodeKJ": {"class": TimerNodeKJ, "name": "Timer Node KJ"},
 
     #instance diffusion
     "CreateInstanceDiffusionTracking": {"class": CreateInstanceDiffusionTracking},
diff --git a/nodes/model_optimization_nodes.py b/nodes/model_optimization_nodes.py
index 5daa5a4..e862631 100644
--- a/nodes/model_optimization_nodes.py
+++ b/nodes/model_optimization_nodes.py
@@ -788,7 +788,7 @@ def teacache_wanvideo_forward_orig(self, x, t, context, clip_fea=None, freqs=Non
         x = self.unpatchify(x, grid_sizes)
         return x
 
-class WanVideoTeaCache:
+class WanVideoTeaCacheKJ:
     @classmethod
     def INPUT_TYPES(s):
         return {
diff --git a/nodes/nodes.py b/nodes/nodes.py
index 38cd54f..3b3580f 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -2601,3 +2601,47 @@ class EmbedND_RifleX(nn.Module):
             dim=-3,
         )
         return emb.unsqueeze(1)
+
+
+class Timer:
+    def __init__(self, name):
+        self.name = name
+        self.start_time = None
+        self.elapsed = 0
+
+class TimerNodeKJ:
+    @classmethod
+    
+    def INPUT_TYPES(s):
+      return {
+        "required": {
+            "any_input": (any, {}),
+            "mode": (["start", "stop"],),
+            "name": ("STRING", {"default": "Timer"}),
+        },
+        "optional": {
+            "timer": ("TIMER",),
+        },
+	}
+
+    RETURN_TYPES = (any, "TIMER", "INT", )
+    RETURN_NAMES = ("any_output", "timer", "time")
+    FUNCTION = "timer"
+    CATEGORY = "KJNodes/misc"
+
+    def timer(self, mode, name, any_input=None, timer=None):
+        if timer is None:
+            if mode == "start":
+                timer = Timer(name=name)            
+                timer.start_time = time.time()
+                return {"ui": {
+                "text": [f"{timer.start_time}"]}, 
+                "result": (any_input, timer, 0) 
+                 }
+        elif mode == "stop" and timer is not None:
+            end_time = time.time()
+            timer.elapsed = int((end_time - timer.start_time) * 1000)
+            timer.start_time = None
+            return (any_input, timer, timer.elapsed)
+
+        
\ No newline at end of file