mirror of
https://github.com/ksyasuda/SubMiner.git
synced 2026-09-30 17:40:53 -07:00
347 lines
12 KiB
TypeScript
347 lines
12 KiB
TypeScript
import { access, mkdtemp, readFile, rm, stat, writeFile } from 'node:fs/promises';
|
|
import { constants } from 'node:fs';
|
|
import { tmpdir } from 'node:os';
|
|
import path from 'node:path';
|
|
import type {
|
|
SubtitleGenerationConfig,
|
|
SubtitleGenerationProgress,
|
|
} from '../../shared/subtitle-generation';
|
|
import { isMissingFile, resolveSubtitleGenerationModel } from './subtitle-generation-models';
|
|
import { runSubtitleGenerationProcess } from './subtitle-generation-process';
|
|
import { publishSubtitleGenerationFile } from './subtitle-generation-files';
|
|
import { formatTimestamp } from './subtitle-generation-srt';
|
|
import { transcribeSubtitleDialogue } from './subtitle-generation-dialogue';
|
|
import {
|
|
loadSubtitleGenerationReference,
|
|
type SubtitleGenerationReference,
|
|
} from './subtitle-generation-reference';
|
|
import {
|
|
requireSubtitleGenerationTools,
|
|
resolveSubtitleGenerationTools,
|
|
} from './subtitle-generation-tools';
|
|
|
|
export {
|
|
downloadSubtitleGenerationModel,
|
|
resolveSubtitleGenerationModel,
|
|
} from './subtitle-generation-models';
|
|
export { resolveSubtitleGenerationTools } from './subtitle-generation-tools';
|
|
|
|
function numericTime(value: unknown): number | undefined {
|
|
if (typeof value !== 'number' && typeof value !== 'string') return undefined;
|
|
const number = Number(value);
|
|
return Number.isFinite(number) ? number : undefined;
|
|
}
|
|
|
|
function parseAudioProbe(raw: string, selectedIndex: number | undefined) {
|
|
const value: unknown = JSON.parse(raw);
|
|
if (
|
|
typeof value !== 'object' ||
|
|
value === null ||
|
|
!('streams' in value) ||
|
|
!Array.isArray(value.streams)
|
|
) {
|
|
throw new Error('ffprobe did not return media streams.');
|
|
}
|
|
const streams = value.streams.flatMap((stream: unknown) => {
|
|
if (
|
|
typeof stream !== 'object' ||
|
|
stream === null ||
|
|
!('codec_type' in stream) ||
|
|
stream.codec_type !== 'audio' ||
|
|
!('index' in stream) ||
|
|
typeof stream.index !== 'number' ||
|
|
!Number.isInteger(stream.index) ||
|
|
stream.index < 0
|
|
)
|
|
return [];
|
|
const tags = 'tags' in stream ? stream.tags : undefined;
|
|
const language =
|
|
typeof tags === 'object' && tags !== null && 'language' in tags ? tags.language : undefined;
|
|
return [
|
|
{
|
|
index: stream.index,
|
|
start: 'start_time' in stream ? numericTime(stream.start_time) : undefined,
|
|
duration: 'duration' in stream ? numericTime(stream.duration) : undefined,
|
|
japanese: language === 'ja' || language === 'jpn',
|
|
},
|
|
];
|
|
});
|
|
const selected =
|
|
selectedIndex === undefined
|
|
? (streams.find((stream) => stream.japanese) ?? streams[0])
|
|
: streams.find((stream) => stream.index === selectedIndex);
|
|
if (!selected)
|
|
throw new Error(
|
|
selectedIndex === undefined
|
|
? 'No audio track found.'
|
|
: `Audio stream ${selectedIndex} was not found.`,
|
|
);
|
|
const format = 'format' in value ? value.format : undefined;
|
|
const formatStart =
|
|
typeof format === 'object' && format !== null && 'start_time' in format
|
|
? (numericTime(format.start_time) ?? 0)
|
|
: 0;
|
|
const duration =
|
|
selected.duration ??
|
|
(typeof format === 'object' && format !== null && 'duration' in format
|
|
? numericTime(format.duration)
|
|
: undefined);
|
|
// mpv rebases media timestamps to the container start. Extraction rebases the selected audio.
|
|
return { index: selected.index, offset: (selected.start ?? formatStart) - formatStart, duration };
|
|
}
|
|
|
|
function shiftSubtitleTimestamps(srt: string, offsetSeconds: number): string {
|
|
let cueCount = 0;
|
|
const result = srt.replace(
|
|
/(\d{2,}):(\d{2}):(\d{2}),(\d{3}) --> (\d{2,}):(\d{2}):(\d{2}),(\d{3})/g,
|
|
(
|
|
_match,
|
|
sh: string,
|
|
sm: string,
|
|
ss: string,
|
|
sms: string,
|
|
eh: string,
|
|
em: string,
|
|
es: string,
|
|
ems: string,
|
|
) => {
|
|
cueCount += 1;
|
|
const start = Number(sh) * 3600000 + Number(sm) * 60000 + Number(ss) * 1000 + Number(sms);
|
|
const end = Number(eh) * 3600000 + Number(em) * 60000 + Number(es) * 1000 + Number(ems);
|
|
return `${formatTimestamp(start + offsetSeconds * 1000)} --> ${formatTimestamp(end + offsetSeconds * 1000)}`;
|
|
},
|
|
);
|
|
if (cueCount === 0)
|
|
throw new Error(
|
|
'Whisper produced no subtitle cues. The audio may contain no recognized speech.',
|
|
);
|
|
return result;
|
|
}
|
|
|
|
async function ensureAvailableOutput(outputPath: string): Promise<void> {
|
|
try {
|
|
await stat(outputPath);
|
|
} catch (error) {
|
|
if (isMissingFile(error)) return;
|
|
throw error;
|
|
}
|
|
throw new Error(`Subtitle output already exists: ${outputPath}`);
|
|
}
|
|
|
|
// Fail before extraction and transcription when the destination cannot take the file.
|
|
export async function ensureWritableDirectory(directory: string): Promise<void> {
|
|
try {
|
|
await access(directory, constants.W_OK | constants.X_OK);
|
|
} catch {
|
|
throw new Error(`Cannot save subtitles: ${directory} is not writable.`);
|
|
}
|
|
}
|
|
|
|
async function writeSubtitles(input: {
|
|
mediaPath: string;
|
|
outputPath?: string;
|
|
contents: string;
|
|
signal?: AbortSignal;
|
|
}): Promise<string> {
|
|
const parsed = path.parse(input.mediaPath);
|
|
const directory = input.outputPath ? path.dirname(path.resolve(input.outputPath)) : parsed.dir;
|
|
const temporaryDirectory = await mkdtemp(path.join(directory, '.subminer-subtitles-'));
|
|
try {
|
|
const staged = path.join(temporaryDirectory, 'subtitles.srt');
|
|
await writeFile(staged, input.contents, { flag: 'wx' });
|
|
for (let suffix = 0; ; suffix += 1) {
|
|
input.signal?.throwIfAborted();
|
|
const destination = input.outputPath
|
|
? path.resolve(input.outputPath)
|
|
: path.join(directory, `${parsed.name}.ja.generated${suffix ? `.${suffix}` : ''}.srt`);
|
|
try {
|
|
await publishSubtitleGenerationFile(staged, destination);
|
|
return destination;
|
|
} catch (error) {
|
|
if (
|
|
!input.outputPath &&
|
|
error instanceof Error &&
|
|
'code' in error &&
|
|
error.code === 'EEXIST'
|
|
)
|
|
continue;
|
|
throw error;
|
|
}
|
|
}
|
|
} finally {
|
|
await rm(temporaryDirectory, { recursive: true, force: true });
|
|
}
|
|
}
|
|
|
|
/** Resolve the model and tools, throwing the message to show when generation cannot run. */
|
|
export async function requireSubtitleGenerationSetup(
|
|
config: SubtitleGenerationConfig,
|
|
modelDirectory: string,
|
|
) {
|
|
const model = await resolveSubtitleGenerationModel(config, modelDirectory);
|
|
if (model.kind === 'missing')
|
|
throw new Error(
|
|
'No Whisper model found. Download a model or configure an existing model path.',
|
|
);
|
|
if (model.kind === 'invalid') throw new Error(model.message);
|
|
const tools = requireSubtitleGenerationTools(await resolveSubtitleGenerationTools(config));
|
|
return { modelPath: model.path, tools };
|
|
}
|
|
|
|
export async function generateJapaneseSubtitles(input: {
|
|
config: SubtitleGenerationConfig;
|
|
modelDirectory: string;
|
|
mediaPath: string;
|
|
audioStreamIndex?: number;
|
|
references?: readonly SubtitleGenerationReference[];
|
|
outputPath?: string;
|
|
/** Parent for the scratch directory holding the extracted WAV; the system temp dir by default. */
|
|
workDirectory?: string;
|
|
onProgress?: (progress: SubtitleGenerationProgress) => void;
|
|
signal?: AbortSignal;
|
|
}): Promise<string> {
|
|
input.signal?.throwIfAborted();
|
|
if (/^[a-z][a-z\d+.-]*:\/\//i.test(input.mediaPath))
|
|
throw new Error('Subtitle generation requires a local media file.');
|
|
const mediaPath = path.resolve(input.mediaPath);
|
|
if (!(await stat(mediaPath)).isFile())
|
|
throw new Error('Subtitle generation requires a local media file.');
|
|
if (input.outputPath) await ensureAvailableOutput(path.resolve(input.outputPath));
|
|
await ensureWritableDirectory(
|
|
input.outputPath ? path.dirname(path.resolve(input.outputPath)) : path.dirname(mediaPath),
|
|
);
|
|
const { modelPath, tools } = await requireSubtitleGenerationSetup(
|
|
input.config,
|
|
input.modelDirectory,
|
|
);
|
|
input.onProgress?.({ stage: 'extract', message: 'Inspecting audio tracks...' });
|
|
const probe = await runSubtitleGenerationProcess({
|
|
command: tools.ffprobe,
|
|
args: [
|
|
'-v',
|
|
'error',
|
|
'-show_entries',
|
|
'stream=index,codec_type,start_time,duration:stream_tags=language:format=start_time,duration',
|
|
'-of',
|
|
'json',
|
|
mediaPath,
|
|
],
|
|
signal: input.signal,
|
|
});
|
|
const audio = parseAudioProbe(probe, input.audioStreamIndex);
|
|
const temporaryDirectory = await mkdtemp(
|
|
path.join(input.workDirectory ?? tmpdir(), 'subminer-whisper-'),
|
|
);
|
|
try {
|
|
const wavPath = path.join(temporaryDirectory, 'audio.wav');
|
|
const subtitleBase = path.join(temporaryDirectory, 'subtitles');
|
|
input.onProgress?.({ stage: 'extract', percent: 0, message: 'Extracting audio...' });
|
|
await runSubtitleGenerationProcess({
|
|
command: tools.ffmpeg,
|
|
args: [
|
|
'-nostdin',
|
|
'-hide_banner',
|
|
'-loglevel',
|
|
'error',
|
|
'-i',
|
|
mediaPath,
|
|
'-map',
|
|
`0:${audio.index}`,
|
|
'-vn',
|
|
'-af',
|
|
'asetpts=PTS-STARTPTS',
|
|
'-ac',
|
|
'1',
|
|
'-ar',
|
|
'16000',
|
|
'-c:a',
|
|
'pcm_s16le',
|
|
'-progress',
|
|
'pipe:1',
|
|
'-nostats',
|
|
wavPath,
|
|
],
|
|
signal: input.signal,
|
|
onLine: (line) => {
|
|
const match = /^out_time_us=(\d+)$/.exec(line);
|
|
if (match && audio.duration && audio.duration > 0) {
|
|
input.onProgress?.({
|
|
stage: 'extract',
|
|
percent: Math.min(100, Math.floor(Number(match[1]) / 10000 / audio.duration)),
|
|
message: 'Extracting audio...',
|
|
});
|
|
}
|
|
},
|
|
});
|
|
input.onProgress?.({
|
|
stage: 'transcribe',
|
|
percent: 0,
|
|
message: 'Generating Japanese subtitles...',
|
|
});
|
|
const referenceStarts = await loadSubtitleGenerationReference({
|
|
references: input.references ?? [],
|
|
mediaPath,
|
|
ffmpegPath: tools.ffmpeg,
|
|
directory: temporaryDirectory,
|
|
audioOffset: audio.offset,
|
|
onProgress: input.onProgress,
|
|
signal: input.signal,
|
|
});
|
|
let srt: string;
|
|
if (tools.vad !== null || referenceStarts.length > 0) {
|
|
srt = await transcribeSubtitleDialogue({
|
|
config: input.config,
|
|
tools,
|
|
referenceStarts,
|
|
modelPath,
|
|
wavPath,
|
|
directory: temporaryDirectory,
|
|
signal: input.signal,
|
|
onProgress: input.onProgress,
|
|
});
|
|
} else {
|
|
await runSubtitleGenerationProcess({
|
|
command: tools.whisper,
|
|
args: [
|
|
'-m',
|
|
modelPath,
|
|
'-f',
|
|
wavPath,
|
|
'-l',
|
|
'ja',
|
|
'-t',
|
|
String(input.config.threads),
|
|
'-osrt',
|
|
'-of',
|
|
subtitleBase,
|
|
'-pp',
|
|
],
|
|
signal: input.signal,
|
|
onLine: (line) => {
|
|
const match = /progress\s*=\s*(\d+(?:\.\d+)?)%/.exec(line);
|
|
if (match)
|
|
input.onProgress?.({
|
|
stage: 'transcribe',
|
|
percent: Math.min(100, Number(match[1])),
|
|
message: 'Generating Japanese subtitles...',
|
|
});
|
|
},
|
|
});
|
|
srt = await readFile(`${subtitleBase}.srt`, 'utf8');
|
|
}
|
|
input.signal?.throwIfAborted();
|
|
input.onProgress?.({ stage: 'write', message: 'Saving Japanese subtitles...' });
|
|
const contents = shiftSubtitleTimestamps(srt, audio.offset);
|
|
const outputPath = await writeSubtitles({
|
|
mediaPath,
|
|
outputPath: input.outputPath,
|
|
contents,
|
|
signal: input.signal,
|
|
});
|
|
input.onProgress?.({ stage: 'write', percent: 100, message: 'Japanese subtitles are ready.' });
|
|
return outputPath;
|
|
} finally {
|
|
await rm(temporaryDirectory, { recursive: true, force: true });
|
|
}
|
|
}
|