Fix inefficient phase unwrapping

Improving_Vocos_vocoder_phase_unwrapping_issue.ipynb

@@ -0,0 +1,127 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "authorship_tag": "ABX9TyOQiYiCFt/wbojXYX3lxI0X",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/roebel/vocos/blob/main/Improving_Vocos_vocoder_phase_unwrapping_issue.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "53P4GXjSkbbS"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "def orig(p):\n",
+        "    x = torch.cos(p)\n",
+        "    y = torch.sin(p)\n",
+        "    phase = torch.atan2(y, x)\n",
+        "    S = torch.exp(phase * 1j)\n",
+        "    return S\n",
+        "\n",
+        "def proposed(p):\n",
+        "    x = torch.cos(p)\n",
+        "    y = torch.sin(p)\n",
+        "    S = x + 1j *y\n",
+        "    return S\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "ta = torch.tile(torch.arange(10000), (10,1))\n",
+        "print(\"compare results:: max diff\", torch.max(torch.abs(orig(ta)-proposed(ta))))"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "tEIYEgb3lJe8",
+        "outputId": "d1644470-dcfd-4878-d13e-70718395ae7e"
+      },
+      "execution_count": 12,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "compare results:: max diff tensor(1.6859e-07)\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "print(\"running original code\")\n",
+        "%timeit orig(ta)\n"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "CgLn8eOLlZO9",
+        "outputId": "65a18b9b-b0db-4bc8-df1f-da8f4646d096"
+      },
+      "execution_count": 6,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "1.98 ms ± 191 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "print(\"running fixed code\")\n",
+        "\n",
+        "%timeit proposed(ta)"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "McRgNqoElpbn",
+        "outputId": "3a6d0b0c-ea07-4ca9-e4f0-f03665829f01"
+      },
+      "execution_count": 10,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "running fixed code\n",
+            "717 µs ± 29.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
+          ]
+        }
+      ]
+    }
+  ]
+}

vocos/heads.py

@@ -54,10 +54,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         mag, p = x.chunk(2, dim=1)
         mag = torch.exp(mag)
         mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
         x = torch.cos(p)
         y = torch.sin(p)
-        phase = torch.atan2(y, x)
-        S = mag * torch.exp(phase * 1j)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value 
+        S = mag * (x + 1j * y)
         audio = self.istft(S)
         return audio