diff --git a/nodes.py b/nodes.py index 8143671..bb0ddcc 100644 --- a/nodes.py +++ b/nodes.py @@ -18,6 +18,10 @@ import torch import comfy.model_management as mm from pathlib import Path + +# Use polygraphy TensorRTX backend +os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1" + from polygraphy import cuda from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife from .trt_utilities import Engine diff --git a/trt_utilities.py b/trt_utilities.py index 143d8e7..eb409e2 100644 --- a/trt_utilities.py +++ b/trt_utilities.py @@ -297,11 +297,10 @@ class Engine: if not enable_all_tactics: config_kwargs["tactic_sources"] = [] - # Load ONNX file then convert the model to fp16 - # Probably might cause issues if onnx file given is quantized or whatever + # Load ONNX file onnx_bytes = onnx.load(onnx_path) onnx_precision = onnx_bytes.graph.input[0].type.tensor_type.elem_type - onnx_simp, check = simplify(float16.convert_float_to_float16(onnx_bytes) if onnx_precision == onnx.TensorProto.FLOAT else onnx_bytes) + onnx_simp, check = simplify(onnx_bytes) if not check: error(f"Simplified ONNX model: {onnx_path} could not be validated") diff --git a/vfi_utilities.py b/vfi_utilities.py index ebc094e..69ef2f0 100644 --- a/vfi_utilities.py +++ b/vfi_utilities.py @@ -11,6 +11,7 @@ import os import torch import typing +import logging import numpy as np from comfy.model_management import soft_empty_cache, get_torch_device @@ -73,6 +74,8 @@ def generate_frames_rife( cache_counter = 0 pbar = ProgressBar(len(frames)) + + logging.info(f"Interpolating frames") for frame_itr in range(len(frames) - 1): @@ -96,13 +99,14 @@ def generate_frames_rife( cache_counter += 1 if cache_counter >= clear_cache_after_n_frames: soft_empty_cache() + logging.info(f"Clearing cache...") cache_counter = 0 pbar.update(1) # Append final frame output_frames[out_len] = frames[-1:] - logger(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}") + logging.info(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}") out_len += 1 # clear cache for courtesy