upgrade Depth

.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
+          asset_content_type: application/x-7z-compressed

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
 

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()
 

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"