class_docstr: SeparateSpeech class
  
  Examples:
      >>> import soundfile
      >>> separate_speech = SeparateSpeech("enh_config.yml", "enh.pth")
      >>> audio, rate = soundfile.read("speech.wav")
      >>> separate_speech(audio)
      [separated_audio1, separated_audio2, ...]
docstr: Inference
  
  Args:
      speech_mix: Input speech data (Batch, Nsamples [, Channels])
      fs: sample rate
  Returns:
      [separated_audio1, separated_audio2, ...]
docstr_tok: ["Inference","Args",":","speech_mix",":","Input","speech","data","(","Batch",",","Nsamples","[",",","Channels","]",")","fs",":","sample","rate","Returns",":","[","separated_audio1",",","separated_audio2",",","...","]"]
code: def __call__(
      self,
      speech_mix: Union[torch.Tensor, np.ndarray],
      fs: int = 8000,
      **kwargs,
  ) -> List[torch.Tensor]:
      """Inference
  
      Args:
          speech_mix: Input speech data (Batch, Nsamples [, Channels])
          fs: sample rate
      Returns:
          [separated_audio1, separated_audio2, ...]
  
      """
      assert check_argument_types()
  
      # Input as audio signal
      if isinstance(speech_mix, np.ndarray):
          speech_mix = torch.as_tensor(speech_mix)
  
      assert speech_mix.dim() > 1, speech_mix.size()
      batch_size = speech_mix.size(0)
      speech_mix = speech_mix.to(getattr(torch, self.dtype))
      # lengths: (B,)
      lengths = speech_mix.new_full(
          [batch_size], dtype=torch.long, fill_value=speech_mix.size(1)
      )
  
      # a. To device
      speech_mix = to_device(speech_mix, device=self.device)
      lengths = to_device(lengths, device=self.device)
  
      ###################################
      # Normalize the signal variance
      if getattr(self.enh_model, "normalize_variance_per_ch", False):
          dim = 1
          mix_std_ = torch.std(speech_mix, dim=dim, keepdim=True)
          speech_mix = speech_mix / mix_std_  # RMS normalization
      elif getattr(self.enh_model, "normalize_variance", False):
          if speech_mix.ndim > 2:
              dim = (1, 2)
          else:
              dim = 1
          mix_std_ = torch.std(speech_mix, dim=dim, keepdim=True)
          speech_mix = speech_mix / mix_std_  # RMS normalization
  
      category = kwargs.get("utt2category", None)
      if (
          self.enh_model.categories
          and category is not None
          and category[0].item() not in self.enh_model.categories
      ):
          raise ValueError(
              f"Category '{category}' is not listed in self.categories"
          )
  
      additional = {}
      if category is not None:
          cat = self.enh_model.categories[category[0].item()]
          print(f"category: {cat}", flush=True)
          if cat.endswith("_reverb"):
              additional["mode"] = "dereverb"
          else:
              additional["mode"] = "no_dereverb"
  
      if self.segmenting and lengths[0] > self.segment_size * fs:
          # Segment-wise speech enhancement/separation
          overlap_length = int(
              np.round(fs * (self.segment_size - self.hop_size))
          )
          num_segments = int(
              np.ceil(
                  (speech_mix.size(1) - overlap_length)
                  / (self.hop_size * fs)
              )
          )
          t = T = int(self.segment_size * fs)
          pad_shape = speech_mix[:, :T].shape
          enh_waves = []
          range_ = trange if self.show_progressbar else range
          for i in range_(num_segments):
              st = int(i * self.hop_size * fs)
              en = st + T
              if en >= lengths[0]:
                  # en - st < T (last segment)
                  en = lengths[0]
                  speech_seg = speech_mix.new_zeros(pad_shape)
                  t = en - st
                  speech_seg[:, :t] = speech_mix[:, st:en]
              else:
                  t = T
                  speech_seg = speech_mix[:, st:en]  # B x T [x C]
  
              lengths_seg = speech_mix.new_full(
                  [batch_size], dtype=torch.long, fill_value=T
              )
              # b. Enhancement/Separation Forward
              feats, f_lens = self.enh_model.encoder(
                  speech_seg, lengths_seg
              )
              if isinstance(self.enh_model, ESPnetDiffusionModel):
                  feats = [self.enh_model.enhance(feats)]
              else:
                  feats, _, _ = self.enh_model.separator(
                      feats, f_lens, additional
                  )
              processed_wav = [
                  self.enh_model.decoder(f, lengths_seg)[0]
                  for f in feats
              ]
              if speech_seg.dim() > 2:
                  # multi-channel speech
                  speech_seg_ = speech_seg[:, self.ref_channel]
              else:
                  speech_seg_ = speech_seg
  
              if self.normalize_segment_scale:
                  # normalize the scale to match the input mixture scale
                  mix_energy = torch.sqrt(
                      torch.mean(
                          speech_seg_[:, :t].pow(2), dim=1, keepdim=True
                      )
                  )
                  enh_energy = torch.sqrt(
                      torch.mean(
                          sum(processed_wav)[:, :t].pow(2),
                          dim=1,
                          keepdim=True,
                      )
                  )
                  processed_wav = [
                      w * (mix_energy / enh_energy)
                      for w in processed_wav
                  ]
              # List[torch.Tensor(num_spk, B, T)]
              enh_waves.append(torch.stack(processed_wav, dim=0))
  
          # c. Stitch the enhanced segments together
          waves = enh_waves[0]
          for i in range(1, num_segments):
              # permutation between separated streams in last and current segments
              perm = self.cal_permumation(
                  waves[:, :, -overlap_length:],
                  enh_waves[i][:, :, :overlap_length],
                  criterion="si_snr",
              )
              # repermute separated streams in current segment
              for batch in range(batch_size):
                  enh_waves[i][:, batch] = enh_waves[i][
                      perm[batch], batch
                  ]
  
              if i == num_segments - 1:
                  enh_waves[i][:, :, t:] = 0
                  enh_waves_res_i = enh_waves[i][:, :, overlap_length:t]
              else:
                  enh_waves_res_i = enh_waves[i][:, :, overlap_length:]
  
              # overlap-and-add (average over the overlapped part)
              waves[:, :, -overlap_length:] = (
                  waves[:, :, -overlap_length:]
                  + enh_waves[i][:, :, :overlap_length]
              ) / 2
              # concatenate the residual parts of the later segment
              waves = torch.cat([waves, enh_waves_res_i], dim=2)
          # ensure the stitched length is same as input
          assert waves.size(2) == speech_mix.size(1), (
              waves.shape,
              speech_mix.shape,
          )
          waves = torch.unbind(waves, dim=0)
      else:
          # b. Enhancement/Separation Forward
          feats, f_lens = self.enh_model.encoder(speech_mix, lengths)
          if isinstance(self.enh_model, ESPnetDiffusionModel):
              feats = [self.enh_model.enhance(feats)]
          else:
              feats, _, _ = self.enh_model.separator(
                  feats, f_lens, additional
              )
          waves = [self.enh_model.decoder(f, lengths)[0] for f in feats]
  
      ###################################
      # De-normalize the signal variance
      if getattr(self.enh_model, "normalize_variance_per_ch", False):
          if mix_std_.ndim > 2:
              mix_std_ = mix_std_[:, :, self.ref_channel]
          waves = [w * mix_std_ for w in waves]
      elif getattr(self.enh_model, "normalize_variance", False):
          if mix_std_.ndim > 2:
              mix_std_ = mix_std_.squeeze(2)
          waves = [w * mix_std_ for w in waves]
  
      assert len(waves) == self.num_spk, len(waves) == self.num_spk
      assert len(waves[0]) == batch_size, (len(waves[0]), batch_size)
      if self.normalize_output_wav:
          waves = [
              (w / abs(w).max(dim=1, keepdim=True)[0] * 0.9)
              .cpu()
              .numpy()
              for w in waves
          ]  # list[(batch, sample)]
      else:
          waves = [w.cpu().numpy() for w in waves]
  
      return waves
  
imports: import argparse
  import logging
  import sys
  from itertools import chain
  from pathlib import Path
  from typing import Any, List, Optional, Sequence, Tuple, Union
  import humanfriendly
  import numpy
  import torch
  import yaml
  from tqdm import trange
  from typeguard import check_argument_types
  from espnet2.enh.diffusion_enh import ESPnetDiffusionModel
  from espnet2.enh.loss.criterions.tf_domain import FrequencyDomainMSE
  from espnet2.enh.loss.criterions.time_domain import SISNRLoss
  from espnet2.enh.loss.wrappers.pit_solver import PITSolver
  from espnet2.fileio.sound_scp import SoundScpWriter
  from espnet2.tasks.enh import EnhancementTask
  from espnet2.tasks.enh_s2t import EnhS2TTask
  from espnet2.torch_utils.device_funcs import to_device
  from espnet2.torch_utils.set_all_random_seed import set_all_random_seed
  from espnet2.train.abs_espnet_model import AbsESPnetModel
  from espnet2.utils import config_argparse
  from espnet2.utils.types import str2bool, str2triple_str, str_or_none
  from espnet.utils.cli_utils import get_commandline_args
label_desc: Use image node_id 2 for calling the SeparateSpeech obj's underlying member method code with example usage: obj.__call__(speech_mix, fs) and returns: waves