From 7bfaa76190231e8d72bf3315661257dc0c2d8c42 Mon Sep 17 00:00:00 2001
From: wangbo <wangbolhr@163.com>
Date: Thu, 15 May 2025 20:58:54 +0800
Subject: [PATCH] fix any

---
 __init__.py |  4 +--
 nodes.py    | 96 +++++++++++++++++++++++++++++++++++++----------------
 2 files changed, 69 insertions(+), 31 deletions(-)

diff --git a/__init__.py b/__init__.py
index dd14aae..be1d75e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,6 +1,6 @@
 WEB_DIRECTORY = "js"
-from .nodes import NODE_CLASS_MAPPINGS
-__all__ = ['NODE_CLASS_MAPPINGS']
+from .nodes import NODE_CLASS_MAPPINGS,NODE_DISPLAY_NAME_MAPPINGS
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
 
 from aiohttp import ClientSession, web
 from server import PromptServer
diff --git a/nodes.py b/nodes.py
index 8da7f29..fceae70 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1,48 +1,86 @@
 import requests
 import io
-import librosa.core as core
+import librosa # Changed from librosa.core
 import torch
+import numpy # For type hinting, though librosa.load returns numpy array
+import warnings
 
 class AudioLoadPath:
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "path": ("STRING", {"default": "X://insert/path/here.mp4"}),
-                              "sample_rate": ("INT", {"default": 22050, "min": 6000, "max": 192000, "step": 1}),
-                              "offset": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1e6, "step": 0.001}),
-                              "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1e6, "step": 0.001})}}
+    def INPUT_TYPES(cls): # Changed s to cls for convention
+        return {
+            "required": {
+                "path": ("STRING", {"default": "X://insert/path/here.mp4"}),
+                "sample_rate": ("INT", {"default": 22050, "min": 6000, "max": 192000, "step": 1}),
+                "offset": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1e6, "step": 0.001}),
+                "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1e6, "step": 0.001})
+            }
+        }
 
-    RETURN_TYPES = ("AUDIO", )
-    CATEGORY = "EasyAI"
+    RETURN_TYPES = ("AUDIO",)
+    CATEGORY = "EasyAI" # Or your preferred category
     FUNCTION = "load"
 
-    def load(self, path: str, sample_rate: int, offset: float, duration: float|None):
+    def load(self, path: str, sample_rate: int, offset: float, duration: float | None):
         if duration == 0.0:
             duration = None
 
-        if path.startswith(('http://', 'https://')):
-            try:
-                response = requests.get(path)
-                response.raise_for_status()
-                audio_data = io.BytesIO(response.content)
+        audio_data_source = None
+        try:
+            if path.startswith(('http://', 'https://')):
+                # For network paths, download and load from memory
+                response = requests.get(path, timeout=10) # Added timeout
+                response.raise_for_status() # Raises an exception for bad status codes
+                audio_data_source = io.BytesIO(response.content)
+            else:
+                # For local paths (absolute or relative)
+                audio_data_source = path
 
-                import warnings
-                with warnings.catch_warnings():
-                    warnings.simplefilter("ignore")
-                    audio, _ = core.load(audio_data, sr=sample_rate, offset=offset, duration=duration)
+            # Use librosa to load audio.
+            # mono=False ensures that the output numpy array is always 2D (channels, samples).
+            # For mono audio, this will be (1, samples).
+            # librosa.load will resample to the target 'sample_rate' if it's provided.
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore") # Suppress librosa warnings if any
+                audio_np, loaded_sr = librosa.load(
+                    audio_data_source,
+                    sr=sample_rate,
+                    offset=offset,
+                    duration=duration,
+                    mono=False # Ensures audio_np is 2D: (channels, samples)
+                )
 
-            except Exception as e:
-                raise Exception(f"加载网络音频失败: {str(e)}")
-        else:
-            audio, _ = core.load(path, sr=sample_rate, offset=offset, duration=duration)
+            # Convert numpy array to PyTorch tensor
+            audio_tensor = torch.from_numpy(audio_np) # Shape: (channels, samples)
 
-            # 修改维度处理方式
-        audio = torch.from_numpy(audio).float()
-        # 确保音频是二维张量 [channels, samples]
-        if audio.dim() == 1:
-            audio = audio.unsqueeze(0)  # 添加channel维度
+            # Add a batch dimension to conform to (batch_size, channels, samples)
+            # Here, batch_size is 1 as we are loading a single audio file.
+            audio_tensor = audio_tensor.unsqueeze(0) # Shape: (1, channels, samples)
 
-        return (audio,)
+            # Prepare the output dictionary for the "AUDIO" type
+            output_audio_dict = {
+                "waveform": audio_tensor,
+                "sample_rate": loaded_sr # Use the actual loaded sample rate (should match input 'sample_rate')
+            }
 
+            # Return as a tuple, as ComfyUI expects
+            return (output_audio_dict,)
+
+        except requests.exceptions.RequestException as e:
+            # Handle network-specific errors
+            raise Exception(f"Failed to load audio from URL: {str(e)}")
+        except FileNotFoundError as e:
+            # Handle local file not found errors
+            raise Exception(f"Audio file not found: {path} - {str(e)}")
+        except Exception as e:
+            # Handle other potential errors (e.g., librosa failing to decode, invalid path)
+            raise Exception(f"Failed to load audio: {str(e)}")
+
+# Node mappings for ComfyUI
 NODE_CLASS_MAPPINGS = {
-    "AudioLoadPath": AudioLoadPath,
+    "AudioLoadPath": AudioLoadPath
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "AudioLoadPath": "Load Audio (Path/URL)"
 }
\ No newline at end of file