From 47f028e0bd5c0a3e3810520a86d792db67006b3c Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 21 Oct 2024 00:29:44 +0300
Subject: [PATCH] Update nodes.py

---
 nodes.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index 5cf0b6f..089daad 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1030,8 +1030,8 @@ class ToraEncodeTrajectory:
             },
         }
 
-    RETURN_TYPES = ("TORAFEATURES",)
-    RETURN_NAMES = ("tora_trajectory",)
+    RETURN_TYPES = ("TORAFEATURES", )
+    RETURN_NAMES = ("tora_trajectory", )
     FUNCTION = "encode"
     CATEGORY = "CogVideoWrapper"
 
@@ -1044,7 +1044,7 @@ class ToraEncodeTrajectory:
         vae = pipeline["pipe"].vae
         vae.enable_slicing()
 
-        canvas_width, canvas_height = 256, 256
+        canvas_width, canvas_height = width, height
         coordinates = json.loads(coordinates.replace("'", '"'))
         coordinates = [(coord['x'], coord['y']) for coord in coordinates]
         
@@ -1081,7 +1081,7 @@ class ToraEncodeTrajectory:
         video_flow_features = traj_extractor(video_flow.to(torch.float32))
         video_flow_features = torch.stack(video_flow_features)
 
-        return (video_flow_features, )   
+        return (video_flow_features,)