c0ffymachyne-ComfyUI_SignalProcessing/core/io.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Author: C0ffymachyne
License: GPLv3
Version: 1.0.0

Description:
    io for audio files and comfy
"""

import torch, torchaudio
from pathlib import Path
from typing import Dict, Tuple, Any
import numpy as np

# this mess needs some cleanup !!


# when we return audio to comfy, turn it to cpu
def audio_to_comfy_3d(
    waveform: torch.Tensor, sample_rate: int, cpu=True, clip: bool = False
) -> Tuple[Dict]:

    # try to fix wrong shape
    if waveform.ndim == 2:
        waveform = waveform.unsqueeze(0)

    if waveform.ndim != 3:
        raise RuntimeError(
            f"returning to_comfy failed {waveform.shape}@{waveform.device}[{waveform.ndim}] has incorrect dimentions, it should be of dimension [batch,channels,audio]"
        )

    if cpu:
        if waveform.device != "cpu":
            waveform = waveform.cpu()

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if clip:
        waveform = np.clip(waveform, -1.0, 1.0)

    return_value = ({"waveform": waveform, "sample_rate": sample_rate},)

    print(
        f"audio_to_comfy_3d '{return_value}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    return return_value


# when we return audio to comfy, turn it to cpu
def audio_to_comfy_3dp1(
    waveform: torch.Tensor, sample_rate: int, value0: Any, cpu=True, clip: bool = False
) -> Tuple[Dict]:

    # try to fix wrong shape
    if waveform.ndim == 2:
        waveform = waveform.unsqueeze(0)

    if waveform.ndim != 3:
        raise RuntimeError(
            f"returning to_comfy failed {waveform.shape}@{waveform.device}dim[{waveform.ndim}] has incorrect dimentions, it should be of dimension [batch,channels,audio]"
        )

    if cpu:
        if waveform.device != "cpu":
            waveform = waveform.cpu()

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if clip:
        waveform = np.clip(waveform, -1.0, 1.0)

    return_value = (
        {"waveform": waveform, "sample_rate": sample_rate},
        value0,
    )

    print(
        f"audio_to_comfy_3d '{return_value}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    return return_value


# when we return audio to comfy, turn it to cpu
def audio_to_comfy_3dp2(
    waveform: torch.Tensor,
    sample_rate: int,
    value0: Any,
    value1: Any,
    cpu=True,
    clip: bool = False,
) -> Tuple[Dict]:

    # try to fix wrong shape
    if waveform.ndim == 2:
        waveform = waveform.unsqueeze(0)

    if waveform.ndim != 3:
        raise RuntimeError(
            f"returning to_comfy failed {waveform.shape}@{waveform.device}dim[{waveform.ndim}] has incorrect dimentions, it should be of dimension [batch,channels,audio]"
        )

    if cpu:
        if waveform.device != "cpu":
            waveform = waveform.cpu()

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if clip:
        waveform = np.clip(waveform, -1.0, 1.0)

    return_value = (
        {"waveform": waveform, "sample_rate": sample_rate},
        value0,
        value1,
    )

    print(
        f"audio_to_comfy_3d '{return_value}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    return return_value


# when we get audio from comfy, load it to gpu
def audio_from_comfy_2d(
    audio: Tuple[Dict], repeat=True, try_gpu=True
) -> Tuple[torch.Tensor, int]:

    if audio is None:
        raise ValueError("Input audio must be provided.")

    waveform = audio.get("waveform")  # [batch, channels, samples]
    sample_rate = audio.get("sample_rate")

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if waveform is None or sample_rate is None:
        raise ValueError("Input audio must contain 'waveform' and 'sample_rate'.")

    if not isinstance(waveform, torch.Tensor):
        raise TypeError("Waveform must be a torch.Tensor.")

    print(
        f"audio_from_comfy_2d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    if waveform.ndim != 3:
        raise ValueError(
            f"Waveform must be a 3D tensor with shape (batch, channels, samples) not {waveform.shape}@{waveform.device}"
        )

    if repeat:
        if waveform.ndim == 1:  # add extra channel in case of mono
            print(f"audio_from_comfy_2d repeat")
            waveform = waveform.unsqueeze(0)  # Add channel dimension
            waveform = waveform.repeat(2, 1)  # copy mono to the new channel

    waveform = waveform.squeeze(0).contiguous()

    if try_gpu:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        waveform = waveform.to(device)

    print(
        f"audio_from_comfy_2d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    return waveform, sample_rate


def audio_from_comfy_3d(
    audio: Tuple[Dict], repeat=True, try_gpu=True
) -> Tuple[torch.Tensor, int]:

    if audio is None:
        raise ValueError("Input audio must be provided.")

    waveform = audio.get("waveform")  # [batch, channels, samples]
    sample_rate = audio.get("sample_rate")

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if waveform is None or sample_rate is None:
        raise ValueError("Input audio must contain 'waveform' and 'sample_rate'.")

    if not isinstance(waveform, torch.Tensor):
        raise TypeError("Waveform must be a torch.Tensor.")

    print(
        f"audio_from_comfy_2d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    if waveform.ndim != 3:
        raise ValueError(
            f"Waveform must be a 3D tensor with shape (batch, channels, samples) not {waveform.shape}@{waveform.device}"
        )

    if repeat:
        if waveform.ndim == 1:  # add extra channel in case of mono
            print(f"audio_from_comfy_2d repeat")
            waveform = waveform.unsqueeze(0)  # Add channel dimension
            waveform = waveform.repeat(2, 1)  # copy mono to the new channel

    # waveform = waveform.squeeze(0).contiguous()

    if try_gpu:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        waveform = waveform.to(device)

    print(
        f"audio_from_comfy_3d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    return waveform, sample_rate


def from_disk_as_dict_3d(
    audio_file: str,
    gain: float = 1.0,
    repeat=True,
    start_time: float = 0.0,
    end_time: float = 0.0,
) -> Dict:
    # all loaded audio comes in 3 dimensions
    # [batch,channels,audio] 3 dimensions

    if not Path(audio_file).exists():
        raise ValueError(f"File {audio_file} Does Not Exist")

    info = torchaudio.info(audio_file)
    sample_rate = info.sample_rate

    # cut sudio segment
    if (start_time > 0.0) and (end_time > start_time):
        frame_offset = int(start_time * sample_rate)  # Start frame
        num_frames = int((end_time - start_time) * sample_rate)  # Total frames to read
        waveform, sample_rate = torchaudio.load(
            audio_file, frame_offset=frame_offset, num_frames=num_frames
        )
    else:
        waveform, sample_rate = torchaudio.load(audio_file)

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)
    waveform = waveform * gain

    if repeat:
        if waveform.ndim == 1:  # add extra channel in case of mono
            waveform = waveform.repeat(2, 1)  # copy mono to the new channel

    # add batch dimension at the front
    waveform = waveform.unsqueeze(0).contiguous()

    # waveform = waveform * gain

    print(
        f"from_disk_as_dict_3d loaded audio file: '{audio_file}' with shape: {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate: {sample_rate} Hz."
    )

    return audio_to_comfy_3d(waveform, sample_rate)


def from_disk_as_raw_3d(
    audio_file: str, repeat: bool = True, try_gpu: bool = False
) -> Tuple[torch.Tensor, int]:
    # all loaded audio comes in 3 dimensions
    # [batch,channels,audio] 3 dimensions

    waveform, sample_rate = torchaudio.load(audio_file)
    info = torchaudio.info(audio_file)
    print(
        f"from_disk_as_raw_3d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate {sample_rate} Hz."
    )

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if repeat:
        if info.num_channels == 1:  # add extra channel in case of mono
            waveform = waveform.repeat(2, 1)  # copy mono to the new channel

    # add batch dimension at the front
    waveform = waveform.unsqueeze(0).contiguous()

    if try_gpu:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        waveform = waveform.to(device)

    print(
        f"from_disk_as_raw_3d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate {sample_rate} Hz."
    )

    return waveform, sample_rate


def from_disk_as_raw_2d(
    audio_file: str, repeat: bool = True, try_gpu: bool = False
) -> Tuple[torch.Tensor, int]:
    # all loaded audio comes in 3 dimensions
    # [batch,channels,audio] 3 dimensions

    waveform, sample_rate = torchaudio.load(audio_file)
    info = torchaudio.info(audio_file)

    print(
        f"from_disk_as_raw_2d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate {sample_rate} Hz."
    )

    if waveform.dtype != torch.float64:
        waveform = waveform.to(dtype=torch.float64)

    if repeat:
        if info.num_channels == 1:  # add extra channel in case of mono
            waveform = waveform.repeat(2, 1)  # copy mono to the new channel
            # waveform = waveform.unsqueeze(0) # Add channel dimension

    if try_gpu:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        waveform = waveform.to(device)

    print(
        f"from_disk_as_raw_2d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
    )

    return waveform, sample_rate