chore: ✨ stash audio experiments

__init__.py

@@ -322,6 +322,14 @@ if hasattr(PromptServer, "instance"):
 
         return await endpoint.do_action(request)
 
+    @PromptServer.instance.routes.get("/mtb/audio")
+    async def get_audio(request):
+        from . import endpoint
+
+        reload(endpoint)
+
+        return await endpoint.get_audio(request)
+
 
 # - WAS Dictionary
 MANIFEST = {

endpoint.py

@@ -3,7 +3,16 @@ import csv
 from aiohttp import web
 
 from .log import mklog
-from .utils import backup_file, here, import_install, reqs_map, run_command, styles_dir
+from .utils import (
+    audioInputDir,
+    backup_file,
+    comfy_dir,
+    here,
+    import_install,
+    reqs_map,
+    run_command,
+    styles_dir,
+)
 
 endlog = mklog("mtb endpoint")
 
@@ -15,6 +24,29 @@ from pathlib import Path
 import_install("requirements")
 
 
+def ACTIONS_loadAudio(args):
+    if not audioInputDir.exists():
+        audioInputDir.mkdir()
+
+    endlog.debug(f"Received Load Audio request for {args}")
+
+    if not args.file:
+        return web.Response(status=400)
+
+    filename = args.filename
+    if not filename:
+        return web.Response(status=400)
+
+    target = audioInputDir / filename
+    if target.exists():
+        target.unlink()
+
+    with target.open("wb") as f:
+        f.write(args.file.read())
+
+    return {"name": filename}
+
+
 def ACTIONS_installDependency(dependency_names=None):
     if dependency_names is None:
         return {"error": "No dependency name provided"}
@@ -88,7 +120,7 @@ def ACTIONS_saveStyle(data):
 
 async def do_action(request) -> web.Response:
     endlog.debug("Init action request")
-    request_data = await request.json()
+    request_data = await request.post()
     name = request_data.get("name")
     args = request_data.get("args")
 
@@ -100,14 +132,33 @@ async def do_action(request) -> web.Response:
     if callable(method):
         result = method(args) if args else method()
         endlog.debug(f"Action result: {result}")
-        return web.json_response({"result": result})
+        return web.json_response({"result": result}, status=200)
 
     available_methods = [
         attr[len("ACTIONS_") :] for attr in globals() if attr.startswith("ACTIONS_")
     ]
 
     return web.json_response(
-        {"error": "Invalid method name.", "available_methods": available_methods}
+        {"error": "Invalid method name.", "available_methods": available_methods},
+        status=400,
+    )
+
+
+async def get_audio(request):
+    name = request.rel_url.query.get("filename")
+    if not name:
+        return web.json_response(
+            {"error": "No filename provided as url query."}, status=400
+        )
+
+    target = audioInputDir / name
+    if not target.exists():
+        return web.json_response(
+            {"error": f"File {name} (in {audioInputDir}) not found..."}, status=404
+        )
+
+    return web.FileResponse(
+        target, headers={"Content-Disposition": f'filename="{name}"'}
     )
 
 

nodes/batch.py

@@ -1,17 +1,20 @@
-import math
-import os
-from pathlib import Path
-from typing import List
+from io import BytesIO
 
 import cv2
-import folder_paths
+import torchaudio
 import numpy as np
 import torch
+from PIL import Image
 
 from ..log import log
 from ..utils import apply_easing, pil2tensor
 from .transform import TransformImage
 
+try:
+    import librosa
+except ImportError:
+    log.warning("librosa not installed. Batch Audio features will not be available.")
+
 
 def hex_to_rgb(hex_color, bgr=False):
     hex_color = hex_color.lstrip("#")
@@ -613,9 +616,57 @@ class BatchShake:
         return (shaken_images, x_translations, y_translations, rotations)
 
 
+class BatchFloatsFromSound:
+    """Extracts a list of floats based on audio frequency band peaks."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "sensitivity": ("FLOAT", {"default": 1.0}),
+                "low_freq": ("FLOAT", {"default": 100.0}),
+                "high_freq": ("FLOAT", {"default": 2000.0}),
+                "hop_length": ("INT", {"default": 512}),
+            },
+        }
+
+    RETURN_TYPES = ("FLOATS",)
+    RETURN_NAMES = ("float_data",)
+    FUNCTION = "process_audio"
+    CATEGORY = "mtb/audio"
+
+    def process_audio(
+        self,
+        audio,
+        sensitivity=1.0,
+        low_freq=100,
+        high_freq=2000,
+        hop_length=512,
+    ):
+        # audio_data, _ = librosa.load(audio_file_path, sr=sample_rate)
+
+        # audio_data_tensor = audio.squeeze(1)  # Remove the channel dimension if present
+        # audio_tensor = audio_data_tensor.float()
+        audio_data = audio.to(device=torchaudio.transforms.Spectrogram().window.device)
+
+        hop_length = 512
+        stft = torchaudio.transforms.Spectrogram()(audio_data)
+        freqs = torchaudio.transforms.FrequencyMasking(low_freq, high_freq)(stft)
+        band_energy = torch.sum(freqs, dim=1)
+
+        min_val = torch.min(band_energy)
+        max_val = torch.max(band_energy)
+        normalized_peaks = (band_energy - min_val) / (max_val - min_val)
+        scaled_peaks = normalized_peaks * sensitivity
+
+        return (scaled_peaks.tolist(),)
+
+
 __nodes__ = [
     BatchFloat,
     Batch2dTransform,
+    BatchFloatsFromSound,
     BatchShape,
     BatchMake,
     BatchFloatAssemble,

nodes/io.py

@@ -1,14 +1,49 @@
-from ..utils import tensor2np, PIL_FILTER_MAP
-import uuid
-import folder_paths
-from ..log import log
-import comfy.model_management as model_management
+import json
 import subprocess
-import torch
+import uuid
 from pathlib import Path
+from typing import List, Optional
+
+import comfy.model_management as model_management
+import folder_paths
 import numpy as np
+import torch
+from comfy.model_management import get_torch_device
 from PIL import Image
-from typing import Optional, List
+
+from ..log import log
+from ..utils import PIL_FILTER_MAP, audioInputDir, tensor2np
+
+try:
+    import librosa
+except ImportError:
+    log.warning("librosa not installed. I/O Audio features will not be available.")
+
+
+class LoadAudio_:
+    """Load an audio file from the input folder (supports upload)"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO_UPLOAD",),
+                "sample_rate": ("INT", {"default": 44100}),
+            }
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "load_audio"
+    CATEGORY = "mtb/audio"
+
+    def load_audio(self, audio: str, sample_rate: int):
+        log.debug(f"Audio file: {audio}")
+        audio_file_path = audioInputDir / audio
+        log.debug(f"Loading audio file: {audio_file_path}")
+        audio_data, _ = librosa.load(audio_file_path.as_posix(), sr=sample_rate)
+        audio_tensor = torch.from_numpy(audio_data).to(get_torch_device())
+        return (audio_tensor.unsqueeze(0).float(),)
 
 
 class ExportWithFfmpeg:
@@ -27,7 +62,8 @@ class ExportWithFfmpeg:
                     ["prores_ks", "libx264", "libx265"],
                     {"default": "prores_ks"},
                 ),
-            }
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
         }
 
     RETURN_TYPES = ("VIDEO",)
@@ -42,9 +78,21 @@ class ExportWithFfmpeg:
         prefix: str,
         format: str,
         codec: str,
+        prompt=None,
+        extra_pnginfo=None,
     ):
+        metadata = {}
         if images.size(0) == 0:
             return ("",)
+
+        if extra_pnginfo is not None:
+            metadata["extra"] = {}
+            for x in extra_pnginfo:
+                metadata["extra"][x] = json.dumps(extra_pnginfo[x])
+
+        if prompt is not None:
+            metadata["prompt"] = json.dumps(prompt)
+
         output_dir = Path(folder_paths.get_output_directory())
         pix_fmt = "rgb48le" if codec == "prores_ks" else "yuv420p"
         file_ext = format
@@ -62,6 +110,15 @@ class ExportWithFfmpeg:
 
         out_path = (output_dir / file_id).as_posix()
 
+        metadata_cmd = []
+
+        if metadata:
+            for k, v in metadata.items():
+                metadata_cmd += [
+                    "-metadata:s:v",
+                    f"{k}='{v if isinstance(v,str) else json.dumps(v)}'",
+                ]
+
         # Prepare the FFmpeg command
         command = [
             "ffmpeg",
@@ -80,6 +137,7 @@ class ExportWithFfmpeg:
             "-",
             "-c:v",
             codec,
+            *metadata_cmd,
             "-r",
             str(fps),
             "-y",
@@ -192,4 +250,4 @@ class SaveGif:
         return {"ui": {"gif": results}}
 
 
-__nodes__ = [SaveGif, ExportWithFfmpeg]
+__nodes__ = [SaveGif, ExportWithFfmpeg, LoadAudio_]

requirements.txt

@@ -5,4 +5,6 @@ requirements-parser
 rembg
 imageio_ffmpeg
 rich
-rich_argparse
+rich_argparse
+librosa
+torchaudio

utils.py

@@ -239,6 +239,7 @@ here = Path(__file__).parent.absolute()
 comfy_dir = Path(folder_paths.base_path)
 models_dir = Path(folder_paths.models_dir)
 styles_dir = comfy_dir / "styles"
+audioInputDir = comfy_dir / "input" / "audio"
 
 # - Construct the path to the font file
 font_path = here / "font.ttf"

web/mtb_widgets.js

@@ -14,7 +14,7 @@ import parseCss from './extern/parse-css.js'
 import * as shared from './comfy_shared.js'
 import { log } from './comfy_shared.js'
 
-const newTypes = [, /*'BOOL'*/ 'COLOR', 'BBOX']
+const newTypes = [, /*'BOOL'*/ 'COLOR', 'BBOX', 'AUDIO_UPLOAD']
 
 const withFont = (ctx, font, cb) => {
   const oldFont = ctx.font
@@ -45,6 +45,54 @@ const calculateTextDimensions = (ctx, value, width, fontSize = 16) => {
   )
   return { textHeight, maxLineWidth }
 }
+function addPlaybackWidget(node, name, url) {
+  let isTick = true
+  const audio = new Audio(url)
+  const slider = node.addWidget(
+    'slider',
+    'loading',
+    0,
+    (v) => {
+      if (!isTick) {
+        audio.currentTime = v
+      }
+      isTick = false
+    },
+    {
+      min: 0,
+      max: 0,
+    }
+  )
+
+  const button = node.addWidget('button', `Play ${name}`, 'play', () => {
+    try {
+      if (audio.paused) {
+        audio.play()
+        button.name = `Pause ${name}`
+      } else {
+        audio.pause()
+        button.name = `Play ${name}`
+      }
+    } catch (error) {
+      alert(error)
+    }
+    app.canvas.setDirty(true)
+  })
+  audio.addEventListener('timeupdate', () => {
+    isTick = true
+    slider.value = audio.currentTime
+    app.canvas.setDirty(true)
+  })
+  audio.addEventListener('ended', () => {
+    button.name = `Play ${name}`
+    app.canvas.setDirty(true)
+  })
+  audio.addEventListener('loadedmetadata', () => {
+    slider.options.max = audio.duration
+    slider.name = `(${audio.duration})`
+    app.canvas.setDirty(true)
+  })
+}
 
 export const MtbWidgets = {
   BBOX: (key, val) => {
@@ -392,6 +440,119 @@ export const MtbWidgets = {
 
     return w
   },
+
+  AUDIO_UPLOAD: function (name, val) {
+    const w = {
+      name,
+      type: 'audio_upload',
+      value: val,
+      draw: function (ctx, node, widgetWidth, widgetY, height) {
+        const [cw, ch] = this.computeSize(widgetWidth)
+        shared.offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, ch)
+      },
+      computeSize: function (width) {
+        if (width) {
+          return [width, 64]
+        }
+        return [128, 128]
+      },
+      onRemoved: function () {
+        if (this.inputEl) {
+          this.inputEl.remove()
+        }
+      },
+    }
+
+    const uploadFile = async (file, node) => {
+      try {
+        const body = new FormData()
+        body.append('name', 'loadAudio')
+        body.append('args', file)
+        const loadAudio = await api.fetchApi('/mtb/actions', {
+          method: 'POST',
+          body,
+        })
+
+        if (loadAudio.status === 200) {
+          const { result } = await loadAudio.json()
+          console.log('received from server', result)
+          console.log(
+            `Getting file /mtb/audio?filename=${encodeURIComponent(
+              result.name
+            )}`
+          )
+
+          w.value = result.name
+          addPlaybackWidget(
+            node,
+            result.name,
+            `/mtb/audio?filename=${encodeURIComponent(result.name)}`
+          )
+        } else {
+          alert(loadAudio.status + ' -' + loadAudio.statusText)
+        }
+        // if (resp.status === 200) {
+        //   const { name } = await resp.json()
+        //   pathWidget.value = name
+        //   addPlaybackWidget(
+        //     node,
+        //     name,
+        //     `/samplediffusion/audio?filename=${encodeURIComponent(name)}`
+        //   )
+        // } else {
+        //   alert(resp.status + ' - ' + resp.statusText)
+        // }
+      } catch (error) {
+        alert(error)
+        throw error
+      }
+    }
+
+    w.inputEl = document.createElement('div')
+    const hidden_input = document.createElement('input')
+    const label = document.createElement('label')
+
+    const uniqueId = 'input_' + Date.now()
+    Object.assign(hidden_input, {
+      type: 'file',
+      accept: 'audio/mpeg,audio/wav,audio/x-wav',
+      id: uniqueId,
+      style: `
+      width: 0.1px;
+      height: 0.1px;
+      opacity: 0;
+      overflow: hidden;
+      position: absolute;
+      z-index: -1;
+
+      `,
+      onchange: async () => {
+        if (hidden_input.files.length) {
+          console.log(hidden_input.files[0])
+          await uploadFile(hidden_input.files[0], this)
+        }
+      },
+    })
+
+    Object.assign(label, {
+      htmlFor: uniqueId,
+    })
+    label.textContent = 'Upload Audio File'
+    label.style = `
+    font-size: 1.25em;
+    font-weight: 700;
+    font-family: monospace;
+    padding:0.5em;
+    border-radius: 5px;
+    color: white;
+    background-color: #1e1e1e;
+    display: inline-block;
+    `
+    document.body.appendChild(w.inputEl)
+    w.inputEl.appendChild(hidden_input)
+    w.inputEl.appendChild(label)
+    return w
+  },
 }
 
 /**
@@ -475,6 +636,20 @@ const mtb_widgets = {
           minHeight: 30,
         }
       },
+      AUDIO_UPLOAD: (node, inputName, inputData, app) => {
+        console.debug('Registering audio')
+        return {
+          widget: node.addCustomWidget(
+            MtbWidgets.AUDIO_UPLOAD.bind(node)(
+              inputName,
+              inputData[1]?.default || ''
+            )
+          ),
+          minWidth: 150,
+          minHeight: 30,
+        }
+      },
+
       // BBOX: (node, inputName, inputData, app) => {
       //     console.debug("Registering bbox")
       //     return {
@@ -887,7 +1062,8 @@ const mtb_widgets = {
 
         break
       }
-      case 'Batch Float Assemble (mtb)': {
+      case 'Batch Float Assemble (mtb)':
+      case 'Plot Batch Float (mtb)': {
         shared.setupDynamicConnections(nodeType, 'floats', 'FLOATS')
         break
       }