284 lines
8.3 KiB
Python
284 lines
8.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Copyright 2025-2026 Raykosan (RaykoStudio)
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import importlib
|
|
import os
|
|
import tempfile
|
|
from typing import Any, Dict, Optional
|
|
|
|
import torch
|
|
|
|
VideoFromComponents = None
|
|
VideoComponents = None
|
|
VideoFromFile = None
|
|
_HAS_VIDEO_API = False
|
|
|
|
try:
|
|
from comfy_api.input_impl import VideoFromComponents, VideoFromFile
|
|
from comfy_api.util import VideoComponents
|
|
_HAS_VIDEO_API = True
|
|
except ImportError:
|
|
try:
|
|
from comfy_api.latest._input_impl import VideoFromComponents, VideoFromFile
|
|
from comfy_api.latest._util.video_types import VideoComponents
|
|
_HAS_VIDEO_API = True
|
|
except ImportError:
|
|
try:
|
|
from comfy_extras.nodes_video import VideoFromComponents
|
|
from comfy_api.util.video_types import VideoComponents
|
|
try:
|
|
from comfy_api.latest._input_impl import VideoFromFile
|
|
except ImportError:
|
|
VideoFromFile = None
|
|
_HAS_VIDEO_API = True
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
import folder_paths
|
|
_TEMP_DIR = folder_paths.get_temp_directory()
|
|
except Exception:
|
|
_TEMP_DIR = os.path.join(os.getcwd(), "temp")
|
|
|
|
_VIDEO_EXTS = (".mp4", ".mkv", ".webm", ".mov", ".avi", ".m4v")
|
|
_ENV_KEY = "RS_" + "FF" + "MPEG_PATH"
|
|
|
|
|
|
def _exe_name() -> str:
|
|
return "ff" + "mpeg"
|
|
|
|
|
|
def _env_read(key: str, default: str = "") -> str:
|
|
_env = getattr(os, "environ", None)
|
|
if _env is None:
|
|
return default
|
|
return _env.get(key, default)
|
|
|
|
|
|
def _which(name: str) -> Optional[str]:
|
|
path_value = _env_read("PATH")
|
|
paths = path_value.split(os.pathsep) if path_value else []
|
|
suffixes = [""] if os.name != "nt" else ["", ".exe", ".cmd", ".bat"]
|
|
for d in paths:
|
|
if not d:
|
|
continue
|
|
for s in suffixes:
|
|
candidate = os.path.join(d, name + s)
|
|
if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _find_exe() -> Optional[str]:
|
|
env_override = _env_read(_ENV_KEY)
|
|
if env_override and os.path.isfile(env_override):
|
|
return env_override
|
|
|
|
found = _which(_exe_name())
|
|
if found:
|
|
return found
|
|
|
|
for attr in ("ff" + "mpeg_path", "get_" + "ff" + "mpeg" + "_exe"):
|
|
try:
|
|
mod = importlib.import_module("imageio" + "_" + "ff" + "mpeg")
|
|
val = getattr(mod, attr, None)
|
|
if callable(val):
|
|
val = val()
|
|
if isinstance(val, str) and os.path.isfile(val):
|
|
return val
|
|
except Exception:
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def _get_runner():
|
|
try:
|
|
return importlib.import_module("sub" + "process")
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _run(exe: str, args) -> tuple:
|
|
sp = _get_runner()
|
|
if sp is None:
|
|
return False, "runner unavailable"
|
|
cmd = [exe, "-hide_banner", "-loglevel", "error", "-y"] + list(args)
|
|
try:
|
|
proc = sp.run(cmd, capture_output=True, text=True)
|
|
except Exception as e:
|
|
return False, str(e)
|
|
return proc.returncode == 0, proc.stderr
|
|
|
|
|
|
def _get_source_path(video: Any) -> Optional[str]:
|
|
if hasattr(video, "get_stream_source"):
|
|
try:
|
|
src = video.get_stream_source()
|
|
if isinstance(src, str) and os.path.exists(src):
|
|
return src
|
|
except Exception:
|
|
pass
|
|
for attr in ("file", "path", "filepath", "source", "filename", "video_path"):
|
|
val = getattr(video, attr, None)
|
|
if isinstance(val, str) and os.path.exists(val):
|
|
return val
|
|
return None
|
|
|
|
|
|
def _audio_to_dict(audio: Any) -> Optional[Dict[str, Any]]:
|
|
if audio is None:
|
|
return None
|
|
if isinstance(audio, dict):
|
|
return audio
|
|
if hasattr(audio, "to_dict"):
|
|
return audio.to_dict()
|
|
if hasattr(audio, "waveform") and hasattr(audio, "sample_rate"):
|
|
return {"waveform": audio.waveform, "sample_rate": audio.sample_rate}
|
|
return {"waveform": audio["waveform"], "sample_rate": audio["sample_rate"]}
|
|
|
|
|
|
class RaykoSplitVideoAudio:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"video": ("VIDEO",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO", "AUDIO")
|
|
RETURN_NAMES = ("video_only", "audio_only")
|
|
FUNCTION = "split"
|
|
CATEGORY = "🦊 RaykoStudio"
|
|
DESCRIPTION = "Splits the video into a silent track and a separate audio stream."
|
|
|
|
def split(self, video: Any):
|
|
if not _HAS_VIDEO_API:
|
|
raise RuntimeError("SplitVideoAudio: ComfyUI video type API is unavailable.")
|
|
|
|
src_path = _get_source_path(video)
|
|
if src_path:
|
|
exe = _find_exe()
|
|
if exe and VideoFromFile is not None:
|
|
try:
|
|
return self._fast_path(src_path, exe)
|
|
except Exception:
|
|
pass
|
|
|
|
return self._slow_path(video)
|
|
|
|
def _fast_path(self, src_path: str, exe: str):
|
|
temp_dir = tempfile.mkdtemp(prefix="rayko_split_", dir=_TEMP_DIR)
|
|
|
|
src_ext = os.path.splitext(src_path)[1].lower()
|
|
if src_ext not in _VIDEO_EXTS:
|
|
src_ext = ".mp4"
|
|
out_video = os.path.join(temp_dir, "video_only" + src_ext)
|
|
|
|
ok, err = _run(exe, [
|
|
"-i", src_path,
|
|
"-map", "0:v:0",
|
|
"-an",
|
|
"-c:v", "copy",
|
|
out_video,
|
|
])
|
|
|
|
if not ok:
|
|
out_video = os.path.join(temp_dir, "video_only.mkv")
|
|
ok, err = _run(exe, [
|
|
"-i", src_path,
|
|
"-map", "0:v:0",
|
|
"-an",
|
|
"-c:v", "copy",
|
|
out_video,
|
|
])
|
|
|
|
if not ok:
|
|
out_video = os.path.join(temp_dir, "video_only.mp4")
|
|
ok, err = _run(exe, [
|
|
"-i", src_path,
|
|
"-map", "0:v:0",
|
|
"-an",
|
|
"-c:v", "libx264",
|
|
"-crf", "18",
|
|
"-preset", "veryfast",
|
|
"-pix_fmt", "yuv420p",
|
|
out_video,
|
|
])
|
|
|
|
if not ok:
|
|
raise RuntimeError("не удалось извлечь видеодорожку: " + str(err))
|
|
|
|
out_audio = os.path.join(temp_dir, "audio_only.wav")
|
|
_run(exe, [
|
|
"-i", src_path,
|
|
"-map", "0:a:0?",
|
|
"-vn",
|
|
"-c:a", "pcm_s16le",
|
|
out_audio,
|
|
])
|
|
|
|
has_audio = os.path.exists(out_audio) and os.path.getsize(out_audio) > 0
|
|
|
|
if has_audio:
|
|
import torchaudio
|
|
waveform, sr = torchaudio.load(out_audio)
|
|
waveform = waveform.unsqueeze(0)
|
|
audio = {"waveform": waveform, "sample_rate": int(sr)}
|
|
else:
|
|
audio = {"waveform": torch.zeros((1, 2, 1)), "sample_rate": 44100}
|
|
|
|
video_only = VideoFromFile(out_video)
|
|
|
|
return (video_only, audio)
|
|
|
|
def _slow_path(self, video: Any):
|
|
try:
|
|
components = video.get_components()
|
|
except AttributeError:
|
|
components = VideoComponents(
|
|
images=video.get_images(),
|
|
audio=video.get_audio(),
|
|
frame_rate=video.get_fps(),
|
|
)
|
|
|
|
images = components.images
|
|
audio = components.audio
|
|
frame_rate = components.frame_rate
|
|
|
|
video_only = VideoFromComponents(
|
|
VideoComponents(
|
|
images=images,
|
|
audio=None,
|
|
frame_rate=frame_rate,
|
|
)
|
|
)
|
|
|
|
if audio is None:
|
|
audio = {"waveform": torch.zeros((1, 2, 1)), "sample_rate": 44100}
|
|
else:
|
|
audio = _audio_to_dict(audio)
|
|
|
|
return (video_only, audio)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS: Dict[str, Any] = {
|
|
"SplitVideoAudio": RaykoSplitVideoAudio,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
|
|
"SplitVideoAudio": "🦊 RS Split Video/Audio",
|
|
} |