Skip to content

Commit 6ecabe6

Browse files
DrewliusDogmaDragon
Drewlius
and
DogmaDragon
authored
Upgrading Stashai to ONNX endpoint (#730)
Co-authored-by: DogmaDragon <103123951+DogmaDragon@users.noreply.github.com>
1 parent 4e6a1ce commit 6ecabe6

2 files changed

Lines changed: 122 additions & 91 deletions

File tree

‎plugins/stashAI/stashai.js‎

Lines changed: 120 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
(function () {
22
"use strict";
33

4-
let STASHMARKER_API_URL = "https://cc1234-stashtag.hf.space/api/predict";
4+
let STASHMARKER_API_BASE = "https://cc1234-stashtag-onnx.hf.space";
55

66
var OPTIONS = [
77
"Anal",
@@ -2241,7 +2241,7 @@
22412241
const [, scene_id] = getScenarioAndID();
22422242
let time;
22432243
let tagId;
2244-
const tagLower = frame.tag.label.toLowerCase();
2244+
const tagLower = frame.tag.label.toLowerCase().replace(/_/g, " ");
22452245

22462246
if (tags[tagLower] === undefined) {
22472247
const tagID = await createTag(tagLower);
@@ -2543,6 +2543,88 @@
25432543
});
25442544
}
25452545

2546+
async function gradioCall(fn_index, image, vtt, threshold, retries = 3) {
2547+
for (let attempt = 0; attempt < retries; attempt++) {
2548+
try {
2549+
return await _gradioCall(fn_index, image, vtt, threshold);
2550+
} catch (err) {
2551+
if (attempt === retries - 1) throw err;
2552+
await new Promise((r) => setTimeout(r, 3000));
2553+
}
2554+
}
2555+
}
2556+
2557+
async function _gradioCall(fn_index, image, vtt, threshold) {
2558+
const session_hash = crypto.randomUUID
2559+
? crypto.randomUUID()
2560+
: "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, (c) => {
2561+
const r = (Math.random() * 16) | 0;
2562+
return (c === "x" ? r : (r & 0x3) | 0x8).toString(16);
2563+
});
2564+
2565+
const queueResponse = await fetch(STASHMARKER_API_BASE + "/gradio_api/queue/join", {
2566+
method: "POST",
2567+
headers: { "Content-Type": "application/json" },
2568+
body: JSON.stringify({
2569+
data: [
2570+
{ url: image, meta: { _type: "gradio.FileData" }, orig_name: "sprite.jpg" },
2571+
vtt,
2572+
threshold,
2573+
],
2574+
fn_index: fn_index,
2575+
session_hash: session_hash,
2576+
}),
2577+
});
2578+
2579+
if (!queueResponse.ok) {
2580+
throw new Error("HTTP " + queueResponse.status);
2581+
}
2582+
2583+
const sseUrl = STASHMARKER_API_BASE + "/gradio_api/queue/data?session_hash=" + session_hash;
2584+
2585+
const controller = new AbortController();
2586+
const timeout = setTimeout(() => controller.abort(), 120000);
2587+
2588+
let text;
2589+
try {
2590+
const resp = await fetch(sseUrl, { signal: controller.signal });
2591+
if (!resp.ok) {
2592+
throw new Error("HTTP " + resp.status);
2593+
}
2594+
text = await resp.text();
2595+
} finally {
2596+
clearTimeout(timeout);
2597+
}
2598+
2599+
let result;
2600+
for (const line of text.split("\n")) {
2601+
if (line.startsWith("data: ")) {
2602+
try {
2603+
const msg = JSON.parse(line.slice(6));
2604+
if (msg.msg === "process_completed") {
2605+
if (msg.success === false) {
2606+
throw new Error(msg.output?.error || "Model inference failed");
2607+
}
2608+
result = [msg.output?.data?.[0]];
2609+
break;
2610+
}
2611+
if (msg.msg === "error") {
2612+
throw new Error(msg.output?.error || "API error");
2613+
}
2614+
} catch (e) {
2615+
if (
2616+
e.message === "Model inference failed" ||
2617+
e.message === "API error"
2618+
)
2619+
throw e;
2620+
}
2621+
}
2622+
}
2623+
2624+
if (result === undefined) throw new Error("No result received");
2625+
return result;
2626+
}
2627+
25462628
function instance$3($$self, $$props, $$invalidate) {
25472629
let { $$slots: slots = {}, $$scope } = $$props;
25482630
validate_slots("MarkerButton", slots, []);
@@ -2569,51 +2651,23 @@
25692651

25702652
let vtt = await download(vtt_url);
25712653

2572-
// query the api with a threshold of 0.4 as we want to do the filtering ourselves
2573-
var data = { data: [image, vtt, 0.4] };
2574-
2575-
fetch(STASHMARKER_API_URL + "_1", {
2576-
method: "POST",
2577-
headers: {
2578-
"Content-Type": "application/json; charset=utf-8",
2579-
},
2580-
body: JSON.stringify(data),
2581-
})
2582-
.then((response) => {
2583-
if (response.status !== 200) {
2584-
$$invalidate(0, (scanner = false));
2585-
alert(
2586-
"Something went wrong. It's likely a server issue, Please try again later."
2587-
);
2588-
return;
2589-
}
2590-
2591-
return response.json();
2592-
})
2593-
.then((data) => {
2594-
$$invalidate(0, (scanner = false));
2595-
let frames = data.data[0];
2596-
$$invalidate(0, (scanner = false));
2597-
2598-
if (frames.length === 0) {
2599-
alert("No tags found");
2600-
return;
2601-
}
2654+
try {
2655+
let result = await gradioCall(1, image, vtt, 0.4);
2656+
let frames = result[0];
26022657

2603-
// find a div with class row
2604-
let row = document.querySelector(".row");
2658+
$$invalidate(0, (scanner = false));
26052659

2606-
new MarkerMatches({ target: row, props: { frames, url } });
2607-
})
2608-
.catch((error) => {
2609-
$$invalidate(0, (scanner = false));
2660+
if (!frames || frames.length === 0) {
2661+
alert("No tags found");
2662+
return;
2663+
}
26102664

2611-
if (error.message === "") {
2612-
alert("Error: Service may be down. please try again later.");
2613-
} else {
2614-
alert("Error: " + error.message);
2615-
}
2616-
});
2665+
let row = document.querySelector(".row");
2666+
new MarkerMatches({ target: row, props: { frames, url } });
2667+
} catch (error) {
2668+
$$invalidate(0, (scanner = false));
2669+
alert("Error: " + (error.message || "Service may be down. Please try again later."));
2670+
}
26172671
}
26182672

26192673
const writable_props = [];
@@ -2630,7 +2684,7 @@
26302684
$$self.$capture_state = () => ({
26312685
getScenarioAndID,
26322686
getUrlSprite,
2633-
STASHMARKER_API_URL,
2687+
STASHMARKER_API_BASE,
26342688
MarkerMatches,
26352689
scanner,
26362690
download,
@@ -3725,11 +3779,12 @@
37253779
let existingTags = await getTagsForScene(scene_id);
37263780

37273781
for (const [tag] of filteredMatches) {
3728-
let tagLower = tag.toLowerCase();
3782+
const tagNormalized = tag.replace(/_/g, " ");
3783+
let tagLower = tagNormalized.toLowerCase();
37293784

37303785
// if tag doesn't exist, create it
37313786
if (tags[tagLower] === undefined) {
3732-
existingTags.push(await createTag(tag));
3787+
existingTags.push(await createTag(tagNormalized));
37333788
} else if (!existingTags.includes(tags[tagLower])) {
37343789
existingTags.push(tags[tagLower]);
37353790
}
@@ -4027,52 +4082,28 @@
40274082
reader.readAsDataURL(vblob);
40284083
});
40294084

4030-
// query the api with a threshold of 0.2 as we want to do the filtering ourselves
4031-
var data = { data: [image, vtt, 0.2] };
4032-
4033-
fetch(STASHMARKER_API_URL, {
4034-
method: "POST",
4035-
headers: {
4036-
"Content-Type": "application/json; charset=utf-8",
4037-
},
4038-
body: JSON.stringify(data),
4039-
})
4040-
.then((response) => {
4041-
if (response.status !== 200) {
4042-
$$invalidate(0, (scanner = false));
4043-
alert(
4044-
"Something went wrong. It's likely a server issue, Please try again later."
4045-
);
4046-
return;
4047-
}
4048-
4049-
return response.json();
4050-
})
4051-
.then((data) => {
4052-
$$invalidate(0, (scanner = false));
4085+
try {
4086+
let result = await gradioCall(0, image, vtt, 0.2);
4087+
let tags = {};
4088+
result.forEach((item) => Object.assign(tags, item));
40534089

4054-
if (data.data[0].length === 0) {
4055-
alert("No tags found");
4056-
return;
4057-
}
4090+
$$invalidate(0, (scanner = false));
40584091

4059-
// grab stash-tag-threshold from local storage or set to default
4060-
let threshold = localStorage.getItem("stash-tag-threshold") || 0.4;
4092+
if (Object.keys(tags).length === 0) {
4093+
alert("No tags found");
4094+
return;
4095+
}
40614096

4062-
new TagMatches({
4063-
target: document.body,
4064-
props: { matches: data.data[0], url, threshold },
4065-
});
4066-
})
4067-
.catch((error) => {
4068-
$$invalidate(0, (scanner = false));
4097+
let threshold = localStorage.getItem("stash-tag-threshold") || 0.4;
40694098

4070-
if (error.message === "") {
4071-
alert("Error: Service may be down. please try again later.");
4072-
} else {
4073-
alert("Error: " + error.message);
4074-
}
4099+
new TagMatches({
4100+
target: document.body,
4101+
props: { matches: tags, url, threshold },
40754102
});
4103+
} catch (error) {
4104+
$$invalidate(0, (scanner = false));
4105+
alert("Error: " + (error.message || "Service may be down. Please try again later."));
4106+
}
40764107
}
40774108

40784109
const writable_props = [];
@@ -4089,7 +4120,7 @@
40894120
$$self.$capture_state = () => ({
40904121
getScenarioAndID,
40914122
getUrlSprite,
4092-
STASHMARKER_API_URL,
4123+
STASHMARKER_API_BASE,
40934124
TagMatches,
40944125
scanner,
40954126
getTags,

‎plugins/stashAI/stashai.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name: Stash AI
22
# requires: CommunityScriptsUILibrary
33
description: Add Tags or Markers to a video from scene sprites using external AI.
4-
version: 1.0.2
4+
version: 1.0.3
55
url: https://discourse.stashapp.cc/t/stash-ai/1392
66
ui:
77
requires:
@@ -12,4 +12,4 @@ ui:
1212
- stashai.css
1313
csp:
1414
connect-src:
15-
- "https://cc1234-stashtag.hf.space"
15+
- "https://cc1234-stashtag-onnx.hf.space"

0 commit comments

Comments
 (0)