Skip to content

Commit 802a9ed

Browse files
Rbell1216Rbell1216
authored
Update worker.js
1 parent 1aebc96 commit 802a9ed

1 file changed

Lines changed: 84 additions & 83 deletions

File tree

‎worker.js‎

Lines changed: 84 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// worker.js - TTS Processing Engine
1+
// worker.js - TTS Processing Engine (WebGPU Only with Retry Logic)
22
import { KokoroTTS } from "./kokoro.js";
33
import { env } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.3.3/dist/transformers.min.js";
44
import { splitTextSmart } from "./semantic-split.js";
@@ -7,29 +7,36 @@ async function detectWebGPU() {
77
try {
88
// Check if WebGPU is supported
99
if (!('gpu' in navigator)) {
10-
console.log("WebGPU not supported in this browser, using WASM");
11-
return false;
10+
console.log("WebGPU not supported in this browser");
11+
throw new Error("WebGPU not supported");
1212
}
1313

1414
const adapter = await navigator.gpu.requestAdapter();
1515
if (!adapter) {
16-
console.log("No WebGPU adapter found, using WASM");
17-
return false;
16+
console.log("No WebGPU adapter found");
17+
throw new Error("No WebGPU adapter found");
1818
}
1919

2020
console.log("WebGPU adapter found:", adapter.info.device || 'Unknown');
2121
return true;
2222
} catch (error) {
23-
console.warn("WebGPU detection failed:", error.message, "using WASM");
24-
return false;
23+
console.error("WebGPU detection failed:", error.message);
24+
throw error;
2525
}
2626
}
2727

28-
const device = await detectWebGPU() ? "webgpu" : "wasm";
28+
let useWebGPU = false;
29+
try {
30+
useWebGPU = await detectWebGPU();
31+
} catch (error) {
32+
console.error("Failed to initialize WebGPU:", error.message);
33+
throw error; // Don't fallback, just fail
34+
}
35+
36+
const device = "webgpu";
2937
self.postMessage({ status: "loading_model_start", device });
3038

3139
let model_id = "onnx-community/Kokoro-82M-v1.0-ONNX";
32-
let useWasmFallback = false;
3340

3441
if (self.location.hostname === "localhost2") {
3542
env.allowLocalModels = true;
@@ -39,32 +46,78 @@ if (self.location.hostname === "localhost2") {
3946
let tts = null;
4047
try {
4148
tts = await KokoroTTS.from_pretrained(model_id, {
42-
dtype: device === "wasm" ? "q8" : "fp32", device,
43-
//dtype: "fp16",
49+
dtype: "fp32",
50+
device: "webgpu",
4451
progress_callback: (progress) => {
4552
self.postMessage({ status: "loading_model_progress", progress });
4653
}
4754
});
4855
} catch (initialError) {
49-
console.warn("Failed to load model with WebGPU, falling back to WASM:", initialError.message);
50-
useWasmFallback = true;
51-
tts = await KokoroTTS.from_pretrained(model_id, {
52-
dtype: "q8",
53-
device: "wasm",
54-
progress_callback: (progress) => {
55-
self.postMessage({ status: "loading_model_progress", progress });
56-
}
57-
});
56+
console.error("Failed to load model with WebGPU:", initialError.message);
57+
throw initialError; // Don't fallback
5858
}
5959

60-
self.postMessage({ status: "loading_model_ready", voices: tts.voices, device: useWasmFallback ? "wasm" : device });
60+
self.postMessage({ status: "loading_model_ready", voices: tts.voices, device: "webgpu" });
6161

62-
// --- THIS IS THE MEMORY-SAFE QUEUE LOGIC FOR SENTENCES WITHIN A CHUNK ---
62+
// --- MEMORY-SAFE QUEUE LOGIC FOR SENTENCES WITHIN A CHUNK ---
6363
let bufferQueueSize = 0;
6464
const MAX_QUEUE_SIZE = 6;
6565
let shouldStop = false;
6666
// --- END QUEUE LOGIC ---
6767

68+
// Retry logic for WebGPU processing
69+
async function processWithRetry(sentence, voice, speed, maxRetries = 3) {
70+
let lastError = null;
71+
72+
for (let attempt = 1; attempt <= maxRetries; attempt++) {
73+
if (shouldStop) break;
74+
75+
try {
76+
// Add timeout for each sentence (25 seconds)
77+
const audioPromise = tts.generate(sentence, { voice, speed });
78+
const timeoutPromise = new Promise((_, reject) => {
79+
setTimeout(() => reject(new Error('Audio generation timeout')), 25000);
80+
});
81+
82+
const audio = await Promise.race([audioPromise, timeoutPromise]);
83+
84+
if (shouldStop) break;
85+
86+
let ab = audio.audio.buffer;
87+
bufferQueueSize++;
88+
self.postMessage({ status: "stream_audio_data", audio: ab }, [ab]);
89+
90+
console.log(`Sentence processed successfully on attempt ${attempt}`);
91+
return true; // Success
92+
93+
} catch (audioError) {
94+
lastError = audioError;
95+
console.warn(`Audio generation attempt ${attempt} failed:`, audioError.message);
96+
97+
// Check for critical errors that shouldn't be retried
98+
if (audioError.message && audioError.message.includes('Session already started')) {
99+
console.error("Session conflict detected, cannot retry");
100+
throw audioError;
101+
}
102+
103+
// If this isn't the last attempt, clear buffers and wait
104+
if (attempt < maxRetries) {
105+
console.log(`Clearing buffers and retrying in 2 seconds (attempt ${attempt + 1}/${maxRetries})`);
106+
107+
// Clear buffer queue to free up memory
108+
bufferQueueSize = 0;
109+
110+
// Wait 2 seconds before retry
111+
await new Promise(resolve => setTimeout(resolve, 2000));
112+
}
113+
}
114+
}
115+
116+
// All retries failed
117+
console.error("All retry attempts failed, skipping sentence:", lastError.message);
118+
throw lastError;
119+
}
120+
68121
self.addEventListener("message", async (e) => {
69122
const { type, text, voice, speed } = e.data;
70123

@@ -83,14 +136,12 @@ self.addEventListener("message", async (e) => {
83136
if (type === "generate" && text) {
84137
shouldStop = false;
85138

86-
// MODERATE APPROACH: Process 2-3 sentences per chunk to balance speed and reliability
87-
let sentences = splitTextSmart(text, 800); // Process 2-3 sentences per chunk
88-
let currentTTS = tts;
89-
let useWasmFallback = false;
139+
// Process 2-3 sentences per chunk for balanced speed and reliability
140+
let sentences = splitTextSmart(text, 800); // Use same chunk size as main.js for consistency
90141

91142
console.log(`Processing ${sentences.length} sentences per chunk for better speed`);
92143

93-
// CRITICAL FIX: Always report the total sentences being processed
144+
// Report the total sentences being processed
94145
self.postMessage({
95146
status: "chunk_progress",
96147
sentencesInChunk: sentences.length,
@@ -118,70 +169,20 @@ self.addEventListener("message", async (e) => {
118169
const sentence = sentences[i];
119170
if (shouldStop) break;
120171

121-
// Moderate delay between sentences (not too long)
172+
// Delay between sentences
122173
if (i > 0) {
123-
await new Promise(resolve => setTimeout(resolve, 800)); // Reduced from 1500 to 800ms
174+
await new Promise(resolve => setTimeout(resolve, 800));
124175
}
125176

126177
try {
127-
// Simple timeout for each sentence (25 seconds)
128-
const audioPromise = currentTTS.generate(sentence, { voice, speed });
129-
const timeoutPromise = new Promise((_, reject) => {
130-
setTimeout(() => reject(new Error('Audio generation timeout')), 25000);
131-
});
132-
133-
const audio = await Promise.race([audioPromise, timeoutPromise]);
134-
135-
if (shouldStop) break;
136-
137-
let ab = audio.audio.buffer;
138-
bufferQueueSize++;
139-
self.postMessage({ status: "stream_audio_data", audio: ab }, [ab]);
178+
// Use retry logic instead of direct generation
179+
await processWithRetry(sentence, voice, speed);
140180

141181
console.log(`Sentence ${i + 1}/${sentences.length} processed successfully`);
142182

143183
} catch (audioError) {
144-
const isWebGPUError = audioError.message &&
145-
(audioError.message.includes('GPUBuffer') ||
146-
audioError.message.includes('webgpu') ||
147-
audioError.name === 'AbortError');
148-
149-
const isSessionError = audioError.message &&
150-
audioError.message.includes('Session already started');
151-
152-
if (isSessionError || (isWebGPUError && !useWasmFallback)) {
153-
console.warn("Session/WebGPU error, falling back to WASM:", audioError.message);
154-
useWasmFallback = true;
155-
try {
156-
if (!tts._wasmFallback) {
157-
console.log("Initializing WASM fallback model...");
158-
tts._wasmFallback = await KokoroTTS.from_pretrained(model_id, {
159-
dtype: "q8",
160-
device: "wasm"
161-
});
162-
}
163-
currentTTS = tts._wasmFallback;
164-
165-
// Shorter delay for WASM retry
166-
await new Promise(resolve => setTimeout(resolve, 1500));
167-
168-
const audio = await currentTTS.generate(sentence, { voice, speed });
169-
let ab = audio.audio.buffer;
170-
bufferQueueSize++;
171-
self.postMessage({ status: "stream_audio_data", audio: ab }, [ab]);
172-
173-
console.log(`Sentence ${i + 1}/${sentences.length} processed with WASM fallback`);
174-
175-
} catch (wasmError) {
176-
console.error("WASM fallback failed, skipping sentence:", wasmError.message);
177-
// Continue to next sentence
178-
}
179-
} else if (audioError.message && audioError.message.includes('timeout')) {
180-
console.warn("Audio generation timeout, skipping sentence:", sentence.substring(0, 100) + "...");
181-
} else {
182-
console.error("Audio generation error, skipping sentence:", audioError.message);
183-
// Continue to next sentence for any other errors
184-
}
184+
// Continue to next sentence if this one fails
185+
console.error(`Skipping sentence ${i + 1}/${sentences.length}:`, audioError.message);
185186
}
186187
}
187188

0 commit comments

Comments
 (0)