verion 0.1
butaixianran
parent
a8621069fa
commit
cb6bdc071b
|
|
@ -0,0 +1 @@
|
||||||
|
scripts/__pycache__/
|
||||||
64
README.md
64
README.md
|
|
@ -1,2 +1,64 @@
|
||||||
# Stable-Diffusion-Webui-Civitai-Helper
|
# Stable-Diffusion-Webui-Civitai-Helper
|
||||||
Stable Diffusion Webui Extension for Civitai, to manage your model much more easily.
|
Stable Diffusion Webui Extension for Civitai, to handle your models much more easily.
|
||||||
|
|
||||||
|
# Feature
|
||||||
|
* Scan all models to download model information and preview images from Civitai.
|
||||||
|
* Modified Build-in "Extra Network" cards, to add following buttons on each card:
|
||||||
|
- 🖼: Modified "replace preview" text into this icon
|
||||||
|
- 🌐: Open this model's Civitai url in a new tab
|
||||||
|
- 💡: Add this model's trigger words to prompt
|
||||||
|
- 🏷: Use this model's preview image's prompt
|
||||||
|
|
||||||
|
# Install
|
||||||
|
Go to SD webui's extension tab, go to `Install from url` sub-tab.
|
||||||
|
Copy this project's url into it, click install.
|
||||||
|
|
||||||
|
Or, just download this project as zip file, unzip it to `Your SD webui/extensions`.
|
||||||
|
|
||||||
|
Then reload UI with "Reload UI" Button in Setting tab.
|
||||||
|
|
||||||
|
Done.
|
||||||
|
|
||||||
|
# How ot use
|
||||||
|
## Scan model
|
||||||
|
Go to extension tab "Civitai Helper". There is a button called "Scan Model".
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
Click it, extension will scan all your models to generate SHA256 hash, and use this hash, to get model information and preview images from civitai.
|
||||||
|
|
||||||
|
For each model, it will create a json file to save all model info from civitai. This model info file will be "Your_model_name.civitai.info" in your model folder.
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
If a model info file is already exists, it will skip this model. If a model can not be find in civitai, it will create an empty model info file, so it won't scan this model twice.
|
||||||
|
|
||||||
|
### Add new models
|
||||||
|
When you have some new models, just click this button again, to get new model's information and preview images.
|
||||||
|
|
||||||
|
## Model Card
|
||||||
|
Open SD webui's build-in "Extra Network" tab, to show model cards.
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
Move your mouse on to the bottom of a model card. It will show 4 icon buttons:
|
||||||
|
- 🖼: Modified "replace preview" text into this icon
|
||||||
|
- 🌐: Open this model's Civitai url in a new tab
|
||||||
|
- 💡: Add this model's trigger words to prompt
|
||||||
|
- 🏷: Use this model's preview image's prompt
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
If you click Refresh Button of extra network, those additional buttons will be removed. You can click `Refresh Civitai Helper` button to bring them back.
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
Enjoy!
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 48 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 18 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 21 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 14 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.8 KiB |
|
|
@ -0,0 +1,312 @@
|
||||||
|
"use strict";
|
||||||
|
|
||||||
|
|
||||||
|
function getActivePrompt() {
|
||||||
|
const currentTab = get_uiCurrentTabContent();
|
||||||
|
switch (currentTab.id) {
|
||||||
|
case "tab_txt2img":
|
||||||
|
return currentTab.querySelector("#txt2img_prompt textarea");
|
||||||
|
case "tab_img2img":
|
||||||
|
return currentTab.querySelector("#img2img_prompt textarea");
|
||||||
|
}
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
function getActiveNegativePrompt() {
|
||||||
|
const currentTab = get_uiCurrentTabContent();
|
||||||
|
switch (currentTab.id) {
|
||||||
|
case "tab_txt2img":
|
||||||
|
return currentTab.querySelector("#txt2img_neg_prompt textarea");
|
||||||
|
case "tab_img2img":
|
||||||
|
return currentTab.querySelector("#img2img_neg_prompt textarea");
|
||||||
|
}
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
//button's click function
|
||||||
|
function open_model_url(model_type, model_name){
|
||||||
|
console.log("start open_model_url");
|
||||||
|
|
||||||
|
//get hidden components of extension
|
||||||
|
let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea");
|
||||||
|
let js_open_url_btn = gradioApp().getElementById("ch_js_open_url_btn");
|
||||||
|
|
||||||
|
|
||||||
|
//msg to python side
|
||||||
|
let msg = {
|
||||||
|
"action": "",
|
||||||
|
"model_type": "",
|
||||||
|
"model_name": "",
|
||||||
|
"prompt": "",
|
||||||
|
"neg_prompt": "",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
msg["action"] = "open_url";
|
||||||
|
msg["model_type"] = model_type;
|
||||||
|
msg["model_name"] = model_name;
|
||||||
|
msg["prompt"] = "";
|
||||||
|
msg["neg_prompt"] = "";
|
||||||
|
|
||||||
|
// fill to msg box
|
||||||
|
js_msg_txtbox.value = JSON.stringify(msg);
|
||||||
|
js_msg_txtbox.dispatchEvent(new Event("input"));
|
||||||
|
|
||||||
|
//click hidden button
|
||||||
|
js_open_url_btn.click();
|
||||||
|
|
||||||
|
console.log("end open_model_url");
|
||||||
|
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
function add_trigger_words(model_type, model_name){
|
||||||
|
console.log("start add_trigger_words");
|
||||||
|
|
||||||
|
//get hidden components of extension
|
||||||
|
let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea");
|
||||||
|
let js_add_trigger_words_btn = gradioApp().getElementById("ch_js_add_trigger_words_btn");
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
//msg to python side
|
||||||
|
let msg = {
|
||||||
|
"action": "",
|
||||||
|
"model_type": "",
|
||||||
|
"model_name": "",
|
||||||
|
"prompt": "",
|
||||||
|
"neg_prompt": "",
|
||||||
|
}
|
||||||
|
|
||||||
|
msg["action"] = "add_trigger_words";
|
||||||
|
msg["model_type"] = model_type;
|
||||||
|
msg["model_name"] = model_name;
|
||||||
|
msg["neg_prompt"] = "";
|
||||||
|
|
||||||
|
// get active prompt
|
||||||
|
let prompt = getActivePrompt();
|
||||||
|
msg["prompt"] = prompt.value;
|
||||||
|
|
||||||
|
// fill to msg box
|
||||||
|
js_msg_txtbox.value = JSON.stringify(msg);
|
||||||
|
js_msg_txtbox.dispatchEvent(new Event("input"));
|
||||||
|
|
||||||
|
//click hidden button
|
||||||
|
js_add_trigger_words_btn.click();
|
||||||
|
|
||||||
|
console.log("end add_trigger_words");
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
function use_preview_prompt(model_type, model_name){
|
||||||
|
console.log("start use_preview_prompt");
|
||||||
|
|
||||||
|
//get hidden components of extension
|
||||||
|
let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea");
|
||||||
|
let js_use_preview_prompt_btn = gradioApp().getElementById("ch_js_use_preview_prompt_btn");
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
//msg to python side
|
||||||
|
let msg = {
|
||||||
|
"action": "",
|
||||||
|
"model_type": "",
|
||||||
|
"model_name": "",
|
||||||
|
"prompt": "",
|
||||||
|
"neg_prompt": "",
|
||||||
|
}
|
||||||
|
|
||||||
|
msg["action"] = "use_preview_prompt";
|
||||||
|
msg["model_type"] = model_type;
|
||||||
|
msg["model_name"] = model_name;
|
||||||
|
|
||||||
|
// get active prompt
|
||||||
|
prompt = getActivePrompt();
|
||||||
|
msg["prompt"] = prompt.value;
|
||||||
|
|
||||||
|
// get active neg prompt
|
||||||
|
let neg_prompt = getActiveNegativePrompt();
|
||||||
|
msg["neg_prompt"] = neg_prompt.value;
|
||||||
|
|
||||||
|
// fill to msg box
|
||||||
|
js_msg_txtbox.value = JSON.stringify(msg);
|
||||||
|
js_msg_txtbox.dispatchEvent(new Event("input"));
|
||||||
|
|
||||||
|
//click hidden button
|
||||||
|
js_use_preview_prompt_btn.click();
|
||||||
|
|
||||||
|
console.log("end use_preview_prompt");
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
onUiLoaded(() => {
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
// get all extra network tabs
|
||||||
|
let tab_prefix_list = ["txt2img", "img2img"];
|
||||||
|
let model_type_list = ["textual_inversion", "hypernetworks", "checkpoints", "lora"];
|
||||||
|
let cardid_suffix = "cards";
|
||||||
|
|
||||||
|
// update extra network tab pages' cards
|
||||||
|
// * replace "replace preview" text button into an icon
|
||||||
|
// * add 3 button to each card:
|
||||||
|
// - open model url 🌐
|
||||||
|
// - add trigger words 💡
|
||||||
|
// - use preview image's prompt 🏷
|
||||||
|
// notice: javascript can not get response from python side
|
||||||
|
// so, these buttons just sent request to python
|
||||||
|
// then, python side gonna open url and update prompt text box, without telling js side.
|
||||||
|
function update_card_for_civitai(){
|
||||||
|
|
||||||
|
|
||||||
|
//change all "replace preview" into an icon
|
||||||
|
let extra_network_id = "";
|
||||||
|
let extra_network_node = null;
|
||||||
|
let addtional_nodes = null;
|
||||||
|
let replace_preview_btn = null;
|
||||||
|
let ul_node = null;
|
||||||
|
let model_name_node = null;
|
||||||
|
let model_name = "";
|
||||||
|
let model_type = "";
|
||||||
|
let cards = null;
|
||||||
|
for (const tab_prefix of tab_prefix_list) {
|
||||||
|
for (const js_model_type of model_type_list) {
|
||||||
|
//get model_type for python side
|
||||||
|
switch (js_model_type) {
|
||||||
|
case "textual_inversion":
|
||||||
|
model_type = "ti";
|
||||||
|
break;
|
||||||
|
case "hypernetworks":
|
||||||
|
model_type = "hyper";
|
||||||
|
break;
|
||||||
|
case "checkpoints":
|
||||||
|
model_type = "ckp";
|
||||||
|
break;
|
||||||
|
case "lora":
|
||||||
|
model_type = "lora";
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!model_type) {
|
||||||
|
console.log("can not get model_type from: " + js_model_type);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
extra_network_id = tab_prefix+"_"+js_model_type+"_"+cardid_suffix;
|
||||||
|
// console.log("searching extra_network_node: " + extra_network_id);
|
||||||
|
extra_network_node = gradioApp().getElementById(extra_network_id);
|
||||||
|
if (!extra_network_node) {
|
||||||
|
console.log("can not find extra_network_node: " + extra_network_id);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// console.log("find extra_network_node: " + extra_network_id);
|
||||||
|
|
||||||
|
// get all card nodes
|
||||||
|
cards = extra_network_node.querySelectorAll(".card");
|
||||||
|
for (let card of cards) {
|
||||||
|
// replace preview text button into icon
|
||||||
|
replace_preview_btn = card.querySelector(".actions .additional a");
|
||||||
|
replace_preview_btn.style.margin = "0px 10px";
|
||||||
|
if (replace_preview_btn) {
|
||||||
|
if (replace_preview_btn.innerHTML == "replace preview") {
|
||||||
|
replace_preview_btn.innerHTML = "🖼";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
//get model name node
|
||||||
|
model_name_node = card.querySelector(".actions .name");
|
||||||
|
if (!model_name_node){
|
||||||
|
console.log("can not find model name node for cards in " + extra_network_id);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// get model name
|
||||||
|
model_name = model_name_node.innerHTML;
|
||||||
|
if (!model_name) {
|
||||||
|
console.log("model name is empty for cards in " + extra_network_id);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
//get ul node, which is the parent of all buttons
|
||||||
|
ul_node = card.querySelector(".actions .additional ul");
|
||||||
|
|
||||||
|
// then we need to add 3 buttons to each ul node:
|
||||||
|
let open_url_node = document.createElement("button");
|
||||||
|
open_url_node.innerHTML = "🌐";
|
||||||
|
open_url_node.title = "Open this model's civitai url";
|
||||||
|
open_url_node.style.margin = "0px 10px";
|
||||||
|
open_url_node.setAttribute("onclick","open_model_url('"+model_type+"', '"+model_name+"')");
|
||||||
|
|
||||||
|
let add_trigger_words_node = document.createElement("button");
|
||||||
|
add_trigger_words_node.innerHTML = "💡";
|
||||||
|
add_trigger_words_node.title = "Add trigger words to prompt";
|
||||||
|
add_trigger_words_node.style.margin = "0px 10px";
|
||||||
|
add_trigger_words_node.setAttribute("onclick","add_trigger_words('"+model_type+"', '"+model_name+"')");
|
||||||
|
|
||||||
|
let use_preview_prompt_node = document.createElement("button");
|
||||||
|
use_preview_prompt_node.innerHTML = "🏷";
|
||||||
|
use_preview_prompt_node.title = "Use promt from preview image";
|
||||||
|
use_preview_prompt_node.style.margin = "0px 10px";
|
||||||
|
use_preview_prompt_node.setAttribute("onclick","use_preview_prompt('"+model_type+"', '"+model_name+"')");
|
||||||
|
|
||||||
|
//add to card
|
||||||
|
ul_node.appendChild(open_url_node);
|
||||||
|
ul_node.appendChild(add_trigger_words_node);
|
||||||
|
ul_node.appendChild(use_preview_prompt_node);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
//run it once
|
||||||
|
update_card_for_civitai();
|
||||||
|
|
||||||
|
let tab_id = ""
|
||||||
|
let extra_tab = null;
|
||||||
|
let extra_toolbar = null;
|
||||||
|
//add refresh button to extra network's toolbar
|
||||||
|
for (let prefix of tab_prefix_list) {
|
||||||
|
tab_id = prefix + "_extra_tabs";
|
||||||
|
extra_tab = gradioApp().getElementById(tab_id);
|
||||||
|
|
||||||
|
//get toolbar
|
||||||
|
extra_toolbar = extra_tab.querySelector("div.flex.border-b-2.flex-wrap");
|
||||||
|
|
||||||
|
if (!extra_toolbar){
|
||||||
|
console.log("can not get extra network toolbar for " + tab_id);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// add refresh button to toolbar
|
||||||
|
let ch_refresh = document.createElement("button");
|
||||||
|
ch_refresh.innerHTML = "Refresh Civitai Helper";
|
||||||
|
ch_refresh.title = "Refresh Civitai Helper's model card buttons";
|
||||||
|
ch_refresh.className = "gr-button gr-button-lg gr-button-secondary";
|
||||||
|
ch_refresh.onclick = update_card_for_civitai;
|
||||||
|
|
||||||
|
extra_toolbar.appendChild(ch_refresh);
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
});
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -0,0 +1,434 @@
|
||||||
|
# -*- coding: UTF-8 -*-
|
||||||
|
# This extension can help you manage your models from civitai. It can download preview, add trigger words, open model page and use the prompt from preview image
|
||||||
|
# repo: https://github.com/butaixianran/
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
import modules.scripts as scripts
|
||||||
|
import gradio as gr
|
||||||
|
import os
|
||||||
|
import webbrowser
|
||||||
|
import requests
|
||||||
|
import random
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import shutil
|
||||||
|
import modules
|
||||||
|
from modules import script_callbacks
|
||||||
|
|
||||||
|
# from modules import images
|
||||||
|
# from modules.processing import process_images, Processed
|
||||||
|
# from modules.processing import Processed
|
||||||
|
# from modules.shared import opts, cmd_opts, state
|
||||||
|
|
||||||
|
|
||||||
|
# init
|
||||||
|
model_folders = {
|
||||||
|
"ti": "embeddings",
|
||||||
|
"hyper": os.path.join("models", "hypernetworks"),
|
||||||
|
"ckp": os.path.join("models", "Stable-diffusion"),
|
||||||
|
"lora": os.path.join("models", "Lora"),
|
||||||
|
}
|
||||||
|
|
||||||
|
model_exts = (".bin", ".pt", ".safetensors", ".ckpt")
|
||||||
|
model_info_exts = ".info"
|
||||||
|
civitai_info_suffix = ".civitai"
|
||||||
|
civitai_hash_api_url = "https://civitai.com/api/v1/model-versions/by-hash/"
|
||||||
|
|
||||||
|
# js action list
|
||||||
|
js_actions = ("open_url", "add_trigger_words", "use_preview_prompt")
|
||||||
|
|
||||||
|
root_path = os.getcwd()
|
||||||
|
|
||||||
|
# print for debugging
|
||||||
|
def printD(msg):
|
||||||
|
print(f"Civitai Helper: {msg}")
|
||||||
|
|
||||||
|
|
||||||
|
def gen_file_sha256(filname):
|
||||||
|
''' calculate file sha256 '''
|
||||||
|
hash_value = ""
|
||||||
|
with open(filname, "rb") as f:
|
||||||
|
sha256obj = hashlib.sha256()
|
||||||
|
sha256obj.update(f.read())
|
||||||
|
hash_value = sha256obj.hexdigest()
|
||||||
|
|
||||||
|
printD("sha256: " + hash_value)
|
||||||
|
return hash_value
|
||||||
|
|
||||||
|
# scan model to generate SHA256, then use this SHA256 to get model info from civitai
|
||||||
|
def scan_model(skip_nsfw_preview):
|
||||||
|
printD("Start scan_model")
|
||||||
|
|
||||||
|
for model_type, model_folder in model_folders.items():
|
||||||
|
folder_path = os.path.join(root_path, model_folder)
|
||||||
|
printD("Scanning path: " + folder_path)
|
||||||
|
for filename in os.listdir(folder_path):
|
||||||
|
# check ext
|
||||||
|
item = os.path.join(folder_path, filename)
|
||||||
|
base, ext = os.path.splitext(item)
|
||||||
|
if ext in model_exts:
|
||||||
|
# find a model
|
||||||
|
# get preview image
|
||||||
|
first_preview = base+".png"
|
||||||
|
sec_preview = base+".preview.png"
|
||||||
|
# get info file
|
||||||
|
info_file = base + civitai_info_suffix + model_info_exts
|
||||||
|
# check info file
|
||||||
|
if not os.path.isfile(info_file):
|
||||||
|
# get model's sha256
|
||||||
|
printD("Generate SHA256 for model: " + filename)
|
||||||
|
hash = gen_file_sha256(item)
|
||||||
|
|
||||||
|
if not hash:
|
||||||
|
printD("failed generate SHA256 for this file.")
|
||||||
|
return
|
||||||
|
|
||||||
|
# use this sha256 to get model info from civitai
|
||||||
|
printD("Request model info from civitai")
|
||||||
|
r = requests.get(civitai_hash_api_url+hash)
|
||||||
|
if not r.ok:
|
||||||
|
if r.status_code == 404:
|
||||||
|
# this is not a civitai model
|
||||||
|
printD("Civitai does not have this model")
|
||||||
|
printD("Write empty model info file")
|
||||||
|
empty_info = {}
|
||||||
|
with open(info_file, 'w') as f:
|
||||||
|
data = json.dumps(empty_info)
|
||||||
|
f.write(data)
|
||||||
|
# go to next file
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
printD("Get errorcode: " + str(r.status_code))
|
||||||
|
printD(r.text)
|
||||||
|
return
|
||||||
|
|
||||||
|
# try to get content
|
||||||
|
content = None
|
||||||
|
try:
|
||||||
|
content = r.json()
|
||||||
|
except Exception as e:
|
||||||
|
printD("Parse response json failed")
|
||||||
|
printD(str(e))
|
||||||
|
printD("response:")
|
||||||
|
printD(r.text)
|
||||||
|
return
|
||||||
|
|
||||||
|
if not content:
|
||||||
|
printD("error, content from civitai is None")
|
||||||
|
return
|
||||||
|
|
||||||
|
# write model info to file
|
||||||
|
printD("Write model info to file: " + info_file)
|
||||||
|
with open(info_file, 'w') as f:
|
||||||
|
data = json.dumps(content)
|
||||||
|
f.write(data)
|
||||||
|
|
||||||
|
# check preview image
|
||||||
|
if not os.path.isfile(sec_preview):
|
||||||
|
# need to download preview image
|
||||||
|
printD("Need preview image for this model")
|
||||||
|
if content["images"]:
|
||||||
|
for img_dict in content["images"]:
|
||||||
|
if "nsfw" in img_dict.keys():
|
||||||
|
if img_dict["nsfw"]:
|
||||||
|
printD("This image is NSFW")
|
||||||
|
if skip_nsfw_preview:
|
||||||
|
printD("Skip NSFW image")
|
||||||
|
continue
|
||||||
|
|
||||||
|
if "url" in img_dict.keys():
|
||||||
|
printD("Sending request for image: " + img_dict["url"])
|
||||||
|
# get image
|
||||||
|
img_r = requests.get(img_dict["url"], stream=True)
|
||||||
|
if not img_r.ok:
|
||||||
|
printD("Get errorcode: " + str(r.status_code))
|
||||||
|
printD(r.text)
|
||||||
|
return
|
||||||
|
|
||||||
|
# write to file
|
||||||
|
with open(sec_preview, 'wb') as f:
|
||||||
|
img_r.raw.decode_content = True
|
||||||
|
shutil.copyfileobj(img_r.raw, f)
|
||||||
|
|
||||||
|
printD("Created Preview image: " + sec_preview)
|
||||||
|
|
||||||
|
# we only need 1 preview image
|
||||||
|
break
|
||||||
|
|
||||||
|
# for testing, we only check 1 model for each type
|
||||||
|
# break
|
||||||
|
|
||||||
|
printD("End scan_model")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# handle request from javascript
|
||||||
|
# parameter: msg - msg from js
|
||||||
|
# return: (action, model_type, model_name, prompt, neg_prompt)
|
||||||
|
def parse_js_msg(msg):
|
||||||
|
printD("Start parse js msg")
|
||||||
|
msg_dict = json.loads(msg)
|
||||||
|
|
||||||
|
if "action" not in msg_dict.keys():
|
||||||
|
printD("Can not find action from js request")
|
||||||
|
return
|
||||||
|
|
||||||
|
if "model_type" not in msg_dict.keys():
|
||||||
|
printD("Can not find model type from js request")
|
||||||
|
return
|
||||||
|
|
||||||
|
if "model_name" not in msg_dict.keys():
|
||||||
|
printD("Can not find model name from js request")
|
||||||
|
return
|
||||||
|
|
||||||
|
if "prompt" not in msg_dict.keys():
|
||||||
|
printD("Can not find prompt from js request")
|
||||||
|
return
|
||||||
|
|
||||||
|
if "neg_prompt" not in msg_dict.keys():
|
||||||
|
printD("Can not find neg_prompt from js request")
|
||||||
|
return
|
||||||
|
|
||||||
|
action = msg_dict["action"]
|
||||||
|
model_type = msg_dict["model_type"]
|
||||||
|
model_name = msg_dict["model_name"]
|
||||||
|
prompt = msg_dict["prompt"]
|
||||||
|
neg_prompt = msg_dict["neg_prompt"]
|
||||||
|
|
||||||
|
if not action:
|
||||||
|
printD("Action from js request is None")
|
||||||
|
return
|
||||||
|
|
||||||
|
if not model_type:
|
||||||
|
printD("model_type from js request is None")
|
||||||
|
return
|
||||||
|
|
||||||
|
if not model_name:
|
||||||
|
printD("model_name from js request is None")
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
if action not in js_actions:
|
||||||
|
printD("Unknow action: " + action)
|
||||||
|
return
|
||||||
|
|
||||||
|
if model_type not in model_folders.keys():
|
||||||
|
printD("Unknow model_type: " + model_type)
|
||||||
|
return
|
||||||
|
|
||||||
|
printD("End parse js msg")
|
||||||
|
|
||||||
|
return (action, model_type, model_name, prompt, neg_prompt)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# get model info file's content by model type and model name
|
||||||
|
# parameter: model_type, model_name
|
||||||
|
# return: model_info_dict
|
||||||
|
def get_model_info(model_type, model_name):
|
||||||
|
if model_type not in model_folders.keys():
|
||||||
|
printD("unknow model type: " + model_type)
|
||||||
|
return None
|
||||||
|
|
||||||
|
model_folder = model_folders[model_type]
|
||||||
|
model_info_filename = model_name + civitai_info_suffix + model_info_exts
|
||||||
|
model_info_filepath = os.path.join(root_path, model_folder, model_info_filename)
|
||||||
|
|
||||||
|
if not os.path.isfile(model_info_filepath):
|
||||||
|
printD("Can not find model info file: " + model_info_filepath)
|
||||||
|
return None
|
||||||
|
|
||||||
|
model_info = None
|
||||||
|
with open(model_info_filepath, 'r') as f:
|
||||||
|
try:
|
||||||
|
model_info = json.load(f)
|
||||||
|
except Exception as e:
|
||||||
|
printD("Selected file is not json: " + model_info_filepath)
|
||||||
|
printD(e)
|
||||||
|
return None
|
||||||
|
|
||||||
|
return model_info
|
||||||
|
|
||||||
|
|
||||||
|
# get civitai's model url and open it in browser
|
||||||
|
# parameter: model_type, model_name
|
||||||
|
def open_model_url(msg):
|
||||||
|
printD("Start open_model_url")
|
||||||
|
|
||||||
|
result = parse_js_msg(msg)
|
||||||
|
if not result:
|
||||||
|
printD("Parsing js ms failed")
|
||||||
|
return
|
||||||
|
|
||||||
|
action, model_type, model_name, prompt, neg_prompt = result
|
||||||
|
|
||||||
|
model_info = get_model_info(model_type, model_name)
|
||||||
|
if not model_info:
|
||||||
|
printD(f"Failed to get model info for {model_type} {model_name}")
|
||||||
|
return
|
||||||
|
|
||||||
|
if "modelId" not in model_info.keys():
|
||||||
|
printD(f"Failed to get model id from info file for {model_type} {model_name}")
|
||||||
|
return
|
||||||
|
|
||||||
|
model_id = model_info["modelId"]
|
||||||
|
if not model_id:
|
||||||
|
printD(f"model id from info file of {model_type} {model_name} is None")
|
||||||
|
return
|
||||||
|
|
||||||
|
url = "https://civitai.com/models/"+str(model_id)
|
||||||
|
|
||||||
|
printD("Open Url: " + url)
|
||||||
|
# open url
|
||||||
|
webbrowser.open_new_tab(url)
|
||||||
|
|
||||||
|
printD("End open_model_url")
|
||||||
|
|
||||||
|
|
||||||
|
# add trigger words to prompt
|
||||||
|
# parameter: model_type, model_name, prompt
|
||||||
|
# return: [new_prompt, new_prompt] - new prompt with trigger words, return twice for txt2img and img2img
|
||||||
|
def add_trigger_words(msg):
|
||||||
|
printD("Start add_trigger_words")
|
||||||
|
|
||||||
|
result = parse_js_msg(msg)
|
||||||
|
if not result:
|
||||||
|
printD("Parsing js ms failed")
|
||||||
|
return
|
||||||
|
|
||||||
|
action, model_type, model_name, prompt, neg_prompt = result
|
||||||
|
|
||||||
|
|
||||||
|
model_info = get_model_info(model_type, model_name)
|
||||||
|
if not model_info:
|
||||||
|
printD(f"Failed to get model info for {model_type} {model_name}")
|
||||||
|
return [prompt, prompt]
|
||||||
|
|
||||||
|
if "trainedWords" not in model_info.keys():
|
||||||
|
printD(f"Failed to get trainedWords from info file for {model_type} {model_name}")
|
||||||
|
return [prompt, prompt]
|
||||||
|
|
||||||
|
trainedWords = model_info["trainedWords"]
|
||||||
|
if not trainedWords:
|
||||||
|
printD(f"No trainedWords from info file for {model_type} {model_name}")
|
||||||
|
return [prompt, prompt]
|
||||||
|
|
||||||
|
if len(trainedWords) == 0:
|
||||||
|
printD(f"trainedWords from info file for {model_type} {model_name} is empty")
|
||||||
|
return [prompt, prompt]
|
||||||
|
|
||||||
|
# get ful trigger words
|
||||||
|
trigger_words = ""
|
||||||
|
for word in trainedWords:
|
||||||
|
trigger_words = trigger_words + word
|
||||||
|
|
||||||
|
new_prompt = prompt + trigger_words
|
||||||
|
printD("trigger_words: " + trigger_words)
|
||||||
|
printD("prompt: " + prompt)
|
||||||
|
printD("new_prompt: " + new_prompt)
|
||||||
|
|
||||||
|
printD("End add_trigger_words")
|
||||||
|
|
||||||
|
# add to prompt
|
||||||
|
return [new_prompt, new_prompt]
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# use preview image's prompt as prompt
|
||||||
|
# parameter: model_type, model_name, prompt, neg_prompt
|
||||||
|
# return: [new_prompt, new_neg_prompt, new_prompt, new_neg_prompt,] - return twice for txt2img and img2img
|
||||||
|
def use_preview_image_prompt(msg):
|
||||||
|
printD("Start use_preview_image_prompt")
|
||||||
|
|
||||||
|
result = parse_js_msg(msg)
|
||||||
|
if not result:
|
||||||
|
printD("Parsing js ms failed")
|
||||||
|
return
|
||||||
|
|
||||||
|
action, model_type, model_name, prompt, neg_prompt = result
|
||||||
|
|
||||||
|
|
||||||
|
model_info = get_model_info(model_type, model_name)
|
||||||
|
if not model_info:
|
||||||
|
printD(f"Failed to get model info for {model_type} {model_name}")
|
||||||
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||||
|
|
||||||
|
if "images" not in model_info.keys():
|
||||||
|
printD(f"Failed to get images from info file for {model_type} {model_name}")
|
||||||
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||||
|
|
||||||
|
images = model_info["images"]
|
||||||
|
if not images:
|
||||||
|
printD(f"No images from info file for {model_type} {model_name}")
|
||||||
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||||
|
|
||||||
|
if len(images) == 0:
|
||||||
|
printD(f"images from info file for {model_type} {model_name} is empty")
|
||||||
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||||
|
|
||||||
|
# get prompt from preview images' meta data
|
||||||
|
preview_prompt = ""
|
||||||
|
preview_neg_prompt = ""
|
||||||
|
for img in images:
|
||||||
|
if "meta" in img.keys():
|
||||||
|
if "prompt" in img["meta"].keys():
|
||||||
|
if img["meta"]["prompt"]:
|
||||||
|
preview_prompt = img["meta"]["prompt"]
|
||||||
|
|
||||||
|
if "negativePrompt" in img["meta"].keys():
|
||||||
|
if img["meta"]["negativePrompt"]:
|
||||||
|
preview_neg_prompt = img["meta"]["negativePrompt"]
|
||||||
|
|
||||||
|
# we only need 1 prompt
|
||||||
|
if preview_prompt:
|
||||||
|
break
|
||||||
|
|
||||||
|
if not preview_prompt:
|
||||||
|
printD(f"There is no prompt of {model_type} {model_name} in its preview image")
|
||||||
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||||
|
|
||||||
|
printD("End use_preview_image_prompt")
|
||||||
|
|
||||||
|
return [preview_prompt, preview_neg_prompt, preview_prompt, preview_neg_prompt]
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def on_ui_tabs():
|
||||||
|
# init
|
||||||
|
|
||||||
|
# get prompt textarea
|
||||||
|
# UI structure
|
||||||
|
# check modules/ui.py, search for txt2img_paste_fields
|
||||||
|
# Negative prompt is the second element
|
||||||
|
txt2img_prompt = modules.ui.txt2img_paste_fields[0][0]
|
||||||
|
txt2img_neg_prompt = modules.ui.txt2img_paste_fields[1][0]
|
||||||
|
img2img_prompt = modules.ui.img2img_paste_fields[0][0]
|
||||||
|
img2img_neg_prompt = modules.ui.img2img_paste_fields[1][0]
|
||||||
|
|
||||||
|
# ====UI====
|
||||||
|
with gr.Blocks(analytics_enabled=False) as civitai_helper:
|
||||||
|
# info
|
||||||
|
gr.Markdown("Civitai Helper's extension tab")
|
||||||
|
|
||||||
|
skip_nsfw_preview_ckb = gr.Checkbox(label="SKip NSFW Preview images", value=False, elem_id="ch_skip_nsfw_preview_ckb")
|
||||||
|
scan_model_btn = gr.Button(value="Scan model", elem_id="ch_scan_model_btn")
|
||||||
|
|
||||||
|
# hidden component for js
|
||||||
|
js_msg_txtbox = gr.Textbox(label="Request Msg From Js", visible=False, lines=1, value="", elem_id="ch_js_msg_txtbox")
|
||||||
|
js_open_url_btn = gr.Button(value="Open Model Url", visible=False, elem_id="ch_js_open_url_btn")
|
||||||
|
js_add_trigger_words_btn = gr.Button(value="Add Trigger Words", visible=False, elem_id="ch_js_add_trigger_words_btn")
|
||||||
|
js_use_preview_prompt_btn = gr.Button(value="Use Prompt from Preview Image", visible=False, elem_id="ch_js_use_preview_prompt_btn")
|
||||||
|
|
||||||
|
# ====events====
|
||||||
|
scan_model_btn.click(scan_model, inputs=[skip_nsfw_preview_ckb])
|
||||||
|
js_open_url_btn.click(open_model_url, inputs=[js_msg_txtbox])
|
||||||
|
js_add_trigger_words_btn.click(add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, img2img_prompt])
|
||||||
|
js_use_preview_prompt_btn.click(use_preview_image_prompt, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, txt2img_neg_prompt, img2img_prompt, img2img_neg_prompt])
|
||||||
|
|
||||||
|
# the third parameter is the element id on html, with a "tab_" as prefix
|
||||||
|
return (civitai_helper , "Civitai Helper", "civitai_helper"),
|
||||||
|
|
||||||
|
script_callbacks.on_ui_tabs(on_ui_tabs)
|
||||||
Loading…
Reference in New Issue