fix cuda graph

__init__.py

@@ -4,11 +4,8 @@ from comfy.model_management import get_torch_device
 from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger
 from .trt_utilities import Engine
 import folder_paths
-<<<<<<< Updated upstream
-=======
 import time
 from polygraphy import cuda
->>>>>>> Stashed changes
 
 ENGINE_DIR = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
 
@@ -21,10 +18,7 @@ class RifeTensorrt:
                 "engine": (os.listdir(ENGINE_DIR),),
                 "clear_cache_after_n_frames": ("INT", {"default": 50, "min": 1, "max": 1000}),
                 "multiplier": ("INT", {"default": 2, "min": 1}),
-<<<<<<< Updated upstream
-=======
                 "use_cuda_graph": ("BOOLEAN", {"default": True}),
->>>>>>> Stashed changes
                 "keep_model_loaded": ("BOOLEAN", {"default": False}),
             },
         }
@@ -32,10 +26,7 @@ class RifeTensorrt:
     RETURN_TYPES = ("IMAGE", )
     FUNCTION = "vfi"
     CATEGORY = "tensorrt"
-<<<<<<< Updated upstream
-=======
     OUTPUT_NODE=True
->>>>>>> Stashed changes
 
     def vfi(
         self,
@@ -43,12 +34,8 @@ class RifeTensorrt:
         engine,
         clear_cache_after_n_frames=50,
         multiplier=2,
-<<<<<<< Updated upstream
-        keep_model_loaded=False
-=======
         use_cuda_graph=True,
         keep_model_loaded=False,
->>>>>>> Stashed changes
     ):
         B, H, W, C = frames.shape
         shape_dict = {
@@ -57,12 +44,7 @@ class RifeTensorrt:
             "output": {"shape": (1, 3, H, W)},
         }
 
-<<<<<<< Updated upstream
-        # cache tensorrt engine in memory
-        cudaStream = torch.cuda.current_stream().cuda_stream
-=======
         cudaStream = cuda.Stream()
->>>>>>> Stashed changes
         engine_path = os.path.join(ENGINE_DIR, engine)
         if (not hasattr(self, 'engine') or self.engine_label != engine):
             self.engine = Engine(engine_path)
@@ -76,16 +58,6 @@ class RifeTensorrt:
         self.engine.allocate_buffers(shape_dict=shape_dict)
 
         frames = preprocess_frames(frames)
-<<<<<<< Updated upstream
-
-        def return_middle_frame(frame_0, frame_1, timestep):
-            timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
-            output = self.engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream)
-
-            result = output['output']
-            return result
-
-=======
 
         def return_middle_frame(frame_0, frame_1, timestep):
             timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
@@ -97,7 +69,6 @@ class RifeTensorrt:
             result = output['output']
             return result
 
->>>>>>> Stashed changes
         result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame)
         out = postprocess_frames(result)