From 6e134b7bd7124e28d212bd15f14285562bf97cea Mon Sep 17 00:00:00 2001
From: Varshith B <varshith15@gmail.com>
Date: Mon, 3 Feb 2025 23:00:16 +0530
Subject: [PATCH] fix: pitch shift example

---
 nodes/audio_utils/__init__.py                 |  3 +-
 nodes/audio_utils/pitch_shift.py              | 32 +++++++++++++++++++
 .../audio-tensor-utils-example-workflow.json  | 24 +++++++++++---
 3 files changed, 53 insertions(+), 6 deletions(-)
 create mode 100644 nodes/audio_utils/pitch_shift.py

diff --git a/nodes/audio_utils/__init__.py b/nodes/audio_utils/__init__.py
index 3015f664..0251eac8 100644
--- a/nodes/audio_utils/__init__.py
+++ b/nodes/audio_utils/__init__.py
@@ -1,6 +1,7 @@
 from .load_audio_tensor import LoadAudioTensor
 from .save_audio_tensor import SaveAudioTensor
+from .pitch_shift import PitchShifter
 
-NODE_CLASS_MAPPINGS = {"LoadAudioTensor": LoadAudioTensor, "SaveAudioTensor": SaveAudioTensor}
+NODE_CLASS_MAPPINGS = {"LoadAudioTensor": LoadAudioTensor, "SaveAudioTensor": SaveAudioTensor, "PitchShifter": PitchShifter}
 
 __all__ = ["NODE_CLASS_MAPPINGS"]
diff --git a/nodes/audio_utils/pitch_shift.py b/nodes/audio_utils/pitch_shift.py
new file mode 100644
index 00000000..4275deeb
--- /dev/null
+++ b/nodes/audio_utils/pitch_shift.py
@@ -0,0 +1,32 @@
+import numpy as np
+import librosa
+
+class PitchShifter:
+    CATEGORY = "audio_utils"
+    RETURN_TYPES = ("WAVEFORM", "INT")
+    FUNCTION = "execute"
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("WAVEFORM",),
+                "sample_rate": ("INT",),
+                "pitch_shift": ("FLOAT", {
+                    "default": 4.0, 
+                    "min": 0.0, 
+                    "max": 12.0, 
+                    "step": 0.5
+                }),
+            }
+        }
+
+    @classmethod
+    def IS_CHANGED(cls):
+        return float("nan")
+
+    def execute(self, audio, sample_rate, pitch_shift):
+        audio_float = audio.astype(np.float32) / 32768.0
+        shifted_audio = librosa.effects.pitch_shift(audio_float, sample_rate, n_steps=pitch_shift)
+        shifted_int16 = np.clip(shifted_audio * 32768.0, -32768, 32767).astype(np.int16)
+        return shifted_int16, sample_rate
diff --git a/workflows/audio-tensor-utils-example-workflow.json b/workflows/audio-tensor-utils-example-workflow.json
index d66d22ca..37609fe9 100644
--- a/workflows/audio-tensor-utils-example-workflow.json
+++ b/workflows/audio-tensor-utils-example-workflow.json
@@ -1,8 +1,7 @@
 {
     "1": {
       "inputs": {
-        "buffer_size": 200.0,
-        "sample_rate": 48000
+        "buffer_size": 500.0
       },
       "class_type": "LoadAudioTensor",
       "_meta": {
@@ -15,12 +14,27 @@
           "1",
           0
         ],
-        "frame_size": 20.0,
-        "sample_rate": 48000
+        "sample_rate": [
+          "1",
+          1
+        ],
+        "pitch_shift": 4.0
+      },
+      "class_type": "PitchShifter",
+      "_meta": {
+        "title": "Pitch Shift"
+      }
+    },
+    "3": {
+      "inputs": {
+        "audio": [
+          "2",
+          0
+        ]
       },
       "class_type": "SaveAudioTensor",
       "_meta": {
         "title": "Save Audio Tensor"
       }
     }
-  }
\ No newline at end of file
+}
\ No newline at end of file