diff --git a/__init__.py b/__init__.py index b9f444a..44dba2c 100644 --- a/__init__.py +++ b/__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)