Merge pull request #2594 from SillyTavern/sd-stop Image generation stop handling

68de3fed59fb991f04b0d1918e1ed0cdb5d95f9a

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
Parents: d5576cd7b78cb84990c4b3c47fbb39c9ed07ce63 d84fef047821d2202cda629735e4ddcc8e9ff7d6
3 files changed, +22 -13Ignore whitespace
public/scripts/extensions/stable-diffusion/button.html+6 -2
@@ -1,4 +1,8 @@
1<div id="sd_gen" class="list-group-item flex-container flexGap5">1<div id="sd_gen" class="list-group-item flex-container flexGap5">
2 <div class="fa-solid fa-paintbrush extensionsMenuExtensionButton" title="Trigger Stable Diffusion" data-i18n="[title]Trigger Stable Diffusion" /></div>2 <div class="fa-solid fa-paintbrush extensionsMenuExtensionButton" title="Trigger Stable Diffusion" data-i18n="[title]Trigger Stable Diffusion"></div>
3 Generate Image3 <span>Generate Image</span>
4</div>
5<div id="sd_stop_gen" class="list-group-item flex-container flexGap5">
6 <div class="fa-solid fa-circle-stop extensionsMenuExtensionButton" title="Abort current image generation task" data-i18n="[title]Abort current image generation task"></div>
7 <span>Stop Image Generation</span>
4</div>8</div>
public/scripts/extensions/stable-diffusion/index.js+13 -8
@@ -37,6 +37,7 @@ const MODULE_NAME = 'sd';
37const UPDATE_INTERVAL = 1000;37const UPDATE_INTERVAL = 1000;
38// This is a 1x1 transparent PNG38// This is a 1x1 transparent PNG
39const PNG_PIXEL = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=';39const PNG_PIXEL = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=';
40const CUSTOM_STOP_EVENT = 'sd_stop_generation';
4041
41const sources = {42const sources = {
42 extras: 'extras',43 extras: 'extras',
@@ -2290,6 +2291,7 @@ async function generatePicture(initiator, args, trigger, message, callback) {
22902291
2291 const dimensions = setTypeSpecificDimensions(generationType);2292 const dimensions = setTypeSpecificDimensions(generationType);
2292 const abortController = new AbortController();2293 const abortController = new AbortController();
2294 const stopButton = document.getElementById('sd_stop_gen');
2293 let negativePromptPrefix = args?.negative || '';2295 let negativePromptPrefix = args?.negative || '';
2294 let imagePath = '';2296 let imagePath = '';
22952297
@@ -2300,9 +2302,8 @@ async function generatePicture(initiator, args, trigger, message, callback) {
2300 const prompt = await getPrompt(generationType, message, trigger, quietPrompt, combineNegatives);2302 const prompt = await getPrompt(generationType, message, trigger, quietPrompt, combineNegatives);
2301 console.log('Processed image prompt:', prompt);2303 console.log('Processed image prompt:', prompt);
23022304
2303 eventSource.once(event_types.GENERATION_STOPPED, stopListener);2305 $(stopButton).show();
2304 context.deactivateSendButtons();2306 eventSource.once(CUSTOM_STOP_EVENT, stopListener);
2305 hideSwipeButtons();
23062307
2307 if (typeof args?._abortController?.addEventListener === 'function') {2308 if (typeof args?._abortController?.addEventListener === 'function') {
2308 args._abortController.addEventListener('abort', stopListener);2309 args._abortController.addEventListener('abort', stopListener);
@@ -2311,13 +2312,13 @@ async function generatePicture(initiator, args, trigger, message, callback) {
2311 imagePath = await sendGenerationRequest(generationType, prompt, negativePromptPrefix, characterName, callback, initiator, abortController.signal);2312 imagePath = await sendGenerationRequest(generationType, prompt, negativePromptPrefix, characterName, callback, initiator, abortController.signal);
2312 } catch (err) {2313 } catch (err) {
2313 console.trace(err);2314 console.trace(err);
2314 throw new Error('SD prompt text generation failed.');2315 toastr.error('SD prompt text generation failed. Reason: ' + err, 'Image Generation');
2316 throw new Error('SD prompt text generation failed. Reason: ' + err);
2315 }2317 }
2316 finally {2318 finally {
2319 $(stopButton).hide();
2317 restoreOriginalDimensions(dimensions);2320 restoreOriginalDimensions(dimensions);
2318 eventSource.removeListener(event_types.GENERATION_STOPPED, stopListener);2321 eventSource.removeListener(CUSTOM_STOP_EVENT, stopListener);
2319 context.activateSendButtons();
2320 showSwipeButtons();
2321 }2322 }
23222323
2323 return imagePath;2324 return imagePath;
@@ -3395,7 +3396,7 @@ async function addSDGenButtons() {
3395 $(document).on('click touchend', function (e) {3396 $(document).on('click touchend', function (e) {
3396 const target = $(e.target);3397 const target = $(e.target);
3397 if (target.is(dropdown) || target.closest(dropdown).length) return;3398 if (target.is(dropdown) || target.closest(dropdown).length) return;
3398 if (target.is(button) && !dropdown.is(':visible') && $('#send_but').is(':visible')) {3399 if ((target.is(button) || target.closest(button).length) && !dropdown.is(':visible') && $('#send_but').is(':visible')) {
3399 e.preventDefault();3400 e.preventDefault();
34003401
3401 dropdown.fadeIn(animation_duration);3402 dropdown.fadeIn(animation_duration);
@@ -3425,6 +3426,10 @@ async function addSDGenButtons() {
3425 generatePicture(initiators.wand, {}, param);3426 generatePicture(initiators.wand, {}, param);
3426 }3427 }
3427 });3428 });
3429
3430 const stopGenButton = $('#sd_stop_gen');
3431 stopGenButton.hide();
3432 stopGenButton.on('click', () => eventSource.emit(CUSTOM_STOP_EVENT));
3428}3433}
34293434
3430function isValidState() {3435function isValidState() {
public/scripts/extensions/translate/buttons.html+3 -3
@@ -1,8 +1,8 @@
1<div id="translate_chat" class="list-group-item flex-container flexGap5">1<div id="translate_chat" class="list-group-item flex-container flexGap5">
2 <div class="fa-solid fa-language extensionsMenuExtensionButton" /></div>2 <div class="fa-solid fa-language extensionsMenuExtensionButton"></div>
3 <span data-i18n="ext_translate_btn_chat">Translate Chat</span>3 <span data-i18n="ext_translate_btn_chat">Translate Chat</span>
4</div>4</div>
5<div id="translate_input_message" class="list-group-item flex-container flexGap5">5<div id="translate_input_message" class="list-group-item flex-container flexGap5">
6 <div class="fa-solid fa-keyboard extensionsMenuExtensionButton" /></div>6 <div class="fa-solid fa-keyboard extensionsMenuExtensionButton"></div>
7 <span data-i18n="ext_translate_btn_input">Translate Input</span>7 <span data-i18n="ext_translate_btn_input">Translate Input</span>
8</div>
8 \ No newline at end of file \ No newline at end of file
8</div>