From 6194a982afcd5c45e00ddff17b9815d898e32882 Mon Sep 17 00:00:00 2001 From: NevermindNilas Date: Mon, 7 Oct 2024 23:12:46 +0300 Subject: [PATCH] upgrade Depth --- .github/workflows/Build-All-Platforms.yaml | 6 +- CHANGELOG.MD | 4 +- src/depth/depth.py | 86 +++++++++++++++++++--- src/downloadModels.py | 2 + 4 files changed, 84 insertions(+), 14 deletions(-) diff --git a/.github/workflows/Build-All-Platforms.yaml b/.github/workflows/Build-All-Platforms.yaml index d847f4d5..c38fb70a 100644 --- a/.github/workflows/Build-All-Platforms.yaml +++ b/.github/workflows/Build-All-Platforms.yaml @@ -16,9 +16,9 @@ jobs: submodules: recursive - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: - python-version: "3.12.4" + python-version: "3.12.7" - name: Extract version id: get_version @@ -190,4 +190,4 @@ jobs: upload_url: ${{ needs.create_release.outputs.upload_url }} asset_path: TAS_${{ needs.create_release.outputs.version }}_Lite_Linux.7z asset_name: TAS_${{ needs.create_release.outputs.version }}_Lite_Linux.7z - asset_content_type: application/x-7z-compressed \ No newline at end of file + asset_content_type: application/x-7z-compressed diff --git a/CHANGELOG.MD b/CHANGELOG.MD index 1fae397c..c106f58a 100644 --- a/CHANGELOG.MD +++ b/CHANGELOG.MD @@ -15,9 +15,11 @@ and this project adheres to [Semantic Versioning](http://semver.org/). #### Improvements - Increase buffer size for encoder from `10 -> 20` +- Depth TensorRT will now dynamically generate ONNX files based on the selected `depth_method`. This will result in significantly higher quality from TRT which should now match the quality of CUDA. This comes with a `25%` tax in performance as well as a `2x` time increase in engine generation but it's well worth it. #### Notes -- `H266` encoding requires FFMPEG version `7.1`, if you have FFMPEG in System Path make sure to upgrade it, otherwise if you downloaded it using TAS, make sure you delete the `ffmpeg` folder located at `%appdata%\Roaming\TheAnimeScripter\` and allow TAS to re-download the latest build. +- `H266` encoding requires FFMPEG version `7.1`, if you have FFMPEG in System Path make sure the versions matches, otherwise if you downloaded it using TAS, make sure you delete the `ffmpeg` folder located at `%appdata%\Roaming\TheAnimeScripter\` and allow TAS to re-download the latest build. +- With the changes made in Depth TensorRT, `.onnx` files will only be temporarely stored, this will result in significantly smaller storage requirements dropping the total `weights` folder size by roughly `1.5GB` ( not accounting for DirectML ). It is heavily encouraged removing the `weights` folder located at `%appdata%\Roaming\TheAnimeScripter\weights` and allowing TAS to regenerate engines. ## [1.9.10] - 2024-09-29 diff --git a/src/depth/depth.py b/src/depth/depth.py index 9519a6f0..5f5d318f 100644 --- a/src/depth/depth.py +++ b/src/depth/depth.py @@ -559,23 +559,25 @@ def __init__( def handleModels(self): self.isCudaAvailable = torch.cuda.is_available() self.device = torch.device("cuda" if self.isCudaAvailable else "cpu") + self.newHeight, self.newWidth = calculateAspectRatio( + self.width, self.height, self.depthQuality + ) + + # Ugly but it works. + self.depth_method = self.depth_method.replace("-tensorrt", "") + self.filename = modelsMap( - model=self.depth_method, modelType="onnx", half=self.half + model=self.depth_method, modelType="pth", half=self.half ) - folderName = self.depth_method.replace("-tensorrt", "-onnx") - if not os.path.exists(os.path.join(weightsDir, folderName, self.filename)): + if not os.path.exists(os.path.join(weightsDir, self.filename, self.filename)): self.modelPath = downloadModels( model=self.depth_method, half=self.half, - modelType="onnx", + modelType="pth", ) else: - self.modelPath = os.path.join(weightsDir, folderName, self.filename) - - self.newHeight, self.newWidth = calculateAspectRatio( - self.width, self.height, self.depthQuality - ) + self.modelPath = os.path.join(weightsDir, self.filename, self.filename) enginePath = self.TensorRTEngineNameHandler( modelPath=self.modelPath, @@ -589,6 +591,65 @@ def handleModels(self): or self.context is None or not os.path.exists(enginePath) ): + from .dpt_v2 import DepthAnythingV2 + + match self.depth_method: + case "small_v2": + method = "vits" + case "base_v2": + method = "vitb" + case "large_v2": + method = "vitl" + case "giant_v2": + raise NotImplementedError("Giant model not available yet") + # method = "vitg" + + model_configs = { + "vits": { + "encoder": "vits", + "features": 64, + "out_channels": [48, 96, 192, 384], + }, + "vitb": { + "encoder": "vitb", + "features": 128, + "out_channels": [96, 192, 384, 768], + }, + "vitl": { + "encoder": "vitl", + "features": 256, + "out_channels": [256, 512, 1024, 1024], + }, + "vitg": { + "encoder": "vitg", + "features": 384, + "out_channels": [1536, 1536, 1536, 1536], + }, + } + + self.model = DepthAnythingV2(**model_configs[method]) + self.model.load_state_dict(torch.load(self.modelPath, map_location="cpu")) + self.model = self.model.to(self.device).eval() + if self.half and self.isCudaAvailable: + self.model = self.model.half() + + dummyInput = torch.zeros( + (1, 3, self.newHeight, self.newWidth), + device=self.device, + dtype=torch.float16 if self.half else torch.float32, + ) + self.modelPath = self.modelPath.replace(".pth", ".onnx") + + torch.onnx.export( + self.model, + dummyInput, + self.modelPath, + opset_version=19, + input_names=["image"], + output_names=["depth"], + dynamic_axes={"image": {0: "batch_size"}, "depth": {0: "batch_size"}}, + ) + self.engine, self.context = self.TensorRTEngineCreator( modelPath=self.modelPath, enginePath=enginePath, @@ -599,6 +660,11 @@ def handleModels(self): inputName=["image"], ) + try: + os.remove(self.modelPath) + except Exception as e: + logging.exception(f"Failed to delete the onnx file, {e}") + self.stream = torch.cuda.Stream() self.dummyInput = torch.zeros( (1, 3, self.newHeight, self.newWidth), @@ -630,7 +696,7 @@ def handleModels(self): self.context.set_input_shape(tensor_name, self.dummyInput.shape) with torch.cuda.stream(self.stream): - for _ in range(10): + for _ in range(3): self.context.execute_async_v3(stream_handle=self.stream.cuda_stream) self.stream.synchronize() diff --git a/src/downloadModels.py b/src/downloadModels.py index 252a6a7d..c50756f1 100644 --- a/src/downloadModels.py +++ b/src/downloadModels.py @@ -621,6 +621,8 @@ def downloadModels( if model.endswith("-tensorrt") or model.endswith("-directml"): if "rife" in model: folderName = model.replace("-tensorrt", "") + elif "small_v2" in model or "base_v2" in model or "large_v2" in model: + folderName = model.replace("-tensorrt", "").replace("-directml", "") else: folderName = model.replace("-tensorrt", "-onnx").replace( "-directml", "-onnx"