Spaces:
Runtime error
Runtime error
File size: 4,665 Bytes
f2dbf59 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import re
import os
import torch
import shutil
import subprocess
import folder_paths
from torch import Tensor
from collections.abc import Mapping
audio_extensions = ['mp3', 'mp4', 'wav', 'ogg']
video_extensions = ['webm', 'mp4', 'mov']
def ffmpeg_suitability(path):
try:
version = subprocess.run([path, "-version"], check=True,
capture_output=True).stdout.decode("utf-8")
except:
return 0
score = 0
#rough layout of the importance of various features
simple_criterion = [("libvpx", 20),("264",10), ("265",3),
("svtav1",5),("libopus", 1)]
for criterion in simple_criterion:
if version.find(criterion[0]) >= 0:
score += criterion[1]
#obtain rough compile year from copyright information
copyright_index = version.find('2000-2')
if copyright_index >= 0:
copyright_year = version[copyright_index+6:copyright_index+9]
if copyright_year.isnumeric():
score += int(copyright_year)
return score
folder_paths.folder_names_and_paths["VHS_video_formats"] = (
[
os.path.join(os.path.dirname(os.path.abspath(__file__)), "video_formats"),
],
[".json"]
)
if "VHS_FORCE_FFMPEG_PATH" in os.environ:
ffmpeg_path = os.environ.get("VHS_FORCE_FFMPEG_PATH")
else:
ffmpeg_paths = []
try:
from imageio_ffmpeg import get_ffmpeg_exe
imageio_ffmpeg_path = get_ffmpeg_exe()
ffmpeg_paths.append(imageio_ffmpeg_path)
except:
if "VHS_USE_IMAGEIO_FFMPEG" in os.environ:
raise
print("Failed to import imageio_ffmpeg")
if "VHS_USE_IMAGEIO_FFMPEG" in os.environ:
ffmpeg_path = imageio_ffmpeg_path
else:
system_ffmpeg = shutil.which("ffmpeg")
if system_ffmpeg is not None:
ffmpeg_paths.append(system_ffmpeg)
if os.path.isfile("ffmpeg"):
ffmpeg_paths.append(os.path.abspath("ffmpeg"))
if os.path.isfile("ffmpeg.exe"):
ffmpeg_paths.append(os.path.abspath("ffmpeg.exe"))
if len(ffmpeg_paths) == 0:
print("No valid ffmpeg found.")
ffmpeg_path = None
elif len(ffmpeg_paths) == 1:
#Evaluation of suitability isn't required, can take sole option
#to reduce startup time
ffmpeg_path = ffmpeg_paths[0]
else:
ffmpeg_path = max(ffmpeg_paths, key=ffmpeg_suitability)
gifski_path = os.environ.get("VHS_GIFSKI", None)
if gifski_path is None:
gifski_path = os.environ.get("JOV_GIFSKI", None)
if gifski_path is None:
gifski_path = shutil.which("gifski")
ytdl_path = os.environ.get("VHS_YTDL", None) or shutil.which('yt-dlp') \
or shutil.which('youtube-dl')
def get_audio(file, start_time=0, duration=0):
args = [ffmpeg_path, "-i", file]
if start_time > 0:
args += ["-ss", str(start_time)]
if duration > 0:
args += ["-t", str(duration)]
try:
#TODO: scan for sample rate and maintain
res = subprocess.run(args + ["-f", "f32le", "-"],
capture_output=True, check=True)
audio = torch.frombuffer(bytearray(res.stdout), dtype=torch.float32)
match = re.search(', (\\d+) Hz, (\\w+), ',res.stderr.decode('utf-8'))
except subprocess.CalledProcessError as e:
raise Exception(f"VHS failed to extract audio from {file}:\n" \
+ e.stderr.decode("utf-8"))
if match:
ar = int(match.group(1))
#NOTE: Just throwing an error for other channel types right now
#Will deal with issues if they come
ac = {"mono": 1, "stereo": 2}[match.group(2)]
else:
ar = 44100
ac = 2
audio = audio.reshape((-1,ac)).transpose(0,1).unsqueeze(0)
return {'waveform': audio, 'sample_rate': ar}
class LazyAudioMap(Mapping):
def __init__(self, file, start_time, duration):
self.file = file
self.start_time=start_time
self.duration=duration
self._dict=None
def __getitem__(self, key):
if self._dict is None:
self._dict = get_audio(self.file, self.start_time, self.duration)
return self._dict[key]
def __iter__(self):
if self._dict is None:
self._dict = get_audio(self.file, self.start_time, self.duration)
return iter(self._dict)
def __len__(self):
if self._dict is None:
self._dict = get_audio(self.file, self.start_time, self.duration)
return len(self._dict)
def lazy_get_audio(file, start_time=0, duration=0):
return LazyAudioMap(file, start_time, duration)
|