v1.3, code refactoring for preparing powerful new features
butaixianran
parent
c50fe40e1e
commit
8c8d8c142d
|
|
@ -1,2 +1,3 @@
|
|||
scripts/__pycache__/
|
||||
scripts/lib/__pycache__/
|
||||
setting.json
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
# Stable-Diffusion-Webui-Civitai-Helper
|
||||
Stable Diffusion Webui 扩展Civitai助手，用于更轻松的管理和使用Civitai模型。
|
||||
|
||||
## 注意
|
||||
最新的v1.3版本，强大，但是是实验性质的。如果碰到问题，可以去Civitai下载1.2.1老版本: [Civitai Url](https://civitai.com/models/16768/civitai-helper-sd-webui-civitai-extension)
|
||||
|
||||
# 功能
|
||||
* 扫描所有模型，从Civitai下载模型信息和预览图
|
||||
* 修改了内置的"Extra Network"模型卡片，每个卡片增加了如下功能按钮:
|
||||
|
|
|
|||
12
README.md
12
README.md
|
|
@ -4,6 +4,9 @@
|
|||
# Stable-Diffusion-Webui-Civitai-Helper
|
||||
Stable Diffusion Webui Extension for Civitai, to handle your models much more easily.
|
||||
|
||||
## Notice
|
||||
Latest v1.3 is powerful but experimental. If you have issue, you can go back to 1.2.1 by download old version from Civitai: [Civitai Url](https://civitai.com/models/16768/civitai-helper-sd-webui-civitai-extension)
|
||||
|
||||
# Features
|
||||
* Scans all models to download model information and preview images from Civitai.
|
||||
* Modified Built-in "Extra Network" cards, to add the following buttons on each card:
|
||||
|
|
@ -78,6 +81,13 @@ Enjoy!
|
|||
* It cannot force a model link to Civitai by model ID for now. This will be added later.
|
||||
|
||||
# Change Log
|
||||
## v1.3
|
||||
* Open url at client side
|
||||
* Link selected model to civitai by url or model id
|
||||
* Save and load extension setting to file
|
||||
* Show button action's output to UI
|
||||
* Code refactoring
|
||||
|
||||
## v1.2.1
|
||||
* Add more error checking to work with different versions of SD webui.
|
||||
|
||||
|
|
@ -91,4 +101,4 @@ Enjoy!
|
|||
* Support subfolders
|
||||
* Check if refresh is needed when clicking "Refresh Civitai Helper"
|
||||
* Add space when adding trigger words
|
||||
* Add memory optimised sha256 as an option
|
||||
* Add memory Optimized sha256 as an option
|
||||
|
|
|
|||
|
|
@ -1,5 +1,55 @@
|
|||
"use strict";
|
||||
|
||||
// get msg from python side from a hidden textbox
|
||||
// normally this is an old msg, need to wait for a new msg
|
||||
function get_ch_py_msg(){
|
||||
console.log("run get_ch_py_msg")
|
||||
const py_msg_txtbox = gradioApp().querySelector("#ch_py_msg_txtbox textarea");
|
||||
if (py_msg_txtbox && py_msg_txtbox.value) {
|
||||
console.log("find py_msg_txtbox");
|
||||
console.log("py_msg_txtbox value: ");
|
||||
console.log(py_msg_txtbox.value)
|
||||
return py_msg_txtbox.value
|
||||
} else {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// get msg from python side from a hidden textbox
|
||||
// if textbox's value is different from old value then it will consider it is a new msg
|
||||
// it will try once in every sencond, until it reach the max try times
|
||||
const get_new_ch_py_msg = (old_value, max_count=3) => new Promise((resolve, reject) => {
|
||||
console.log("run get_new_ch_py_msg")
|
||||
|
||||
let count = 0;
|
||||
let new_msg = "";
|
||||
let find_msg = false;
|
||||
const interval = setInterval(() => {
|
||||
const py_msg_txtbox = gradioApp().querySelector("#ch_py_msg_txtbox textarea");
|
||||
count++;
|
||||
|
||||
if (py_msg_txtbox && py_msg_txtbox.value) {
|
||||
console.log("find py_msg_txtbox");
|
||||
console.log("py_msg_txtbox value: ");
|
||||
console.log(py_msg_txtbox.value)
|
||||
|
||||
new_msg = py_msg_txtbox.value
|
||||
if (new_msg != old_value) {
|
||||
find_msg=true
|
||||
}
|
||||
}
|
||||
|
||||
if (find_msg) {
|
||||
resolve(new_msg);
|
||||
clearInterval(interval);
|
||||
} else if (count > max_count) {
|
||||
reject('');
|
||||
clearInterval(interval);
|
||||
}
|
||||
|
||||
}, 1000);
|
||||
})
|
||||
|
||||
function getActivePrompt() {
|
||||
const currentTab = get_uiCurrentTabContent();
|
||||
|
|
@ -25,7 +75,7 @@ function getActiveNegativePrompt() {
|
|||
|
||||
|
||||
//button's click function
|
||||
function open_model_url(event, model_type, search_term){
|
||||
async function open_model_url(event, model_type, search_term){
|
||||
console.log("start open_model_url");
|
||||
|
||||
//get hidden components of extension
|
||||
|
|
@ -53,15 +103,39 @@ function open_model_url(event, model_type, search_term){
|
|||
js_msg_txtbox.value = JSON.stringify(msg);
|
||||
js_msg_txtbox.dispatchEvent(new Event("input"));
|
||||
|
||||
//get old py msg
|
||||
let py_msg = get_ch_py_msg();
|
||||
|
||||
//click hidden button
|
||||
js_open_url_btn.click();
|
||||
|
||||
console.log("end open_model_url");
|
||||
|
||||
|
||||
// stop parent event
|
||||
event.stopPropagation()
|
||||
event.preventDefault()
|
||||
|
||||
//check response msg from python
|
||||
let new_py_msg = await get_new_ch_py_msg("");
|
||||
console.log("new_py_msg:");
|
||||
console.log(new_py_msg);
|
||||
|
||||
//check msg
|
||||
if (new_py_msg) {
|
||||
let py_msg_json = JSON.parse(new_py_msg);
|
||||
//check for url
|
||||
if (py_msg_json && py_msg_json.content) {
|
||||
if (py_msg_json.content.url) {
|
||||
window.open(py_msg_json.content.url, "_blank");
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
|
||||
console.log("end open_model_url");
|
||||
|
||||
|
||||
}
|
||||
|
||||
function add_trigger_words(event, model_type, search_term){
|
||||
|
|
|
|||
|
|
@ -1,526 +0,0 @@
|
|||
# -*- 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 re
|
||||
import modules
|
||||
from modules import script_callbacks
|
||||
from modules import shared
|
||||
|
||||
# from modules import images
|
||||
# from modules.processing import process_images, Processed
|
||||
# from modules.processing import Processed
|
||||
# from modules.shared import opts, cmd_opts, state
|
||||
|
||||
# print for debugging
|
||||
def printD(msg):
|
||||
print(f"Civitai Helper: {msg}")
|
||||
|
||||
|
||||
# this is the default root path
|
||||
root_path = os.getcwd()
|
||||
|
||||
# if command line arguement is used to change model folder,
|
||||
# then model folder is in absolute path, not based on this root path anymore.
|
||||
# so to make extension work with those absolute model folder paths, model folder also need to be in absolute path
|
||||
model_folders = {
|
||||
"ti": os.path.join(root_path, "embeddings"),
|
||||
"hyper": os.path.join(root_path, "models", "hypernetworks"),
|
||||
"ckp": os.path.join(root_path, "models", "Stable-diffusion"),
|
||||
"lora": os.path.join(root_path, "models", "Lora"),
|
||||
}
|
||||
|
||||
|
||||
# get cusomter model path
|
||||
# will be modified when refactoring
|
||||
if shared.cmd_opts.embeddings_dir and os.path.isdir(shared.cmd_opts.embeddings_dir):
|
||||
model_folders["ti"] = shared.cmd_opts.embeddings_dir
|
||||
|
||||
if shared.cmd_opts.hypernetwork_dir and os.path.isdir(shared.cmd_opts.hypernetwork_dir):
|
||||
model_folders["hyper"] = shared.cmd_opts.hypernetwork_dir
|
||||
|
||||
if shared.cmd_opts.ckpt_dir and os.path.isdir(shared.cmd_opts.ckpt_dir):
|
||||
model_folders["ckp"] = shared.cmd_opts.ckpt_dir
|
||||
|
||||
if shared.cmd_opts.lora_dir and os.path.isdir(shared.cmd_opts.lora_dir):
|
||||
model_folders["lora"] = shared.cmd_opts.lora_dir
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
|
||||
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
|
||||
|
||||
def gen_file_sha256_low_memory(filname):
|
||||
printD("Using Memory Optimised SHA256")
|
||||
hash_sha256 = hashlib.sha256()
|
||||
with open(filname, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_sha256.update(chunk)
|
||||
|
||||
hash_value = hash_sha256.hexdigest()
|
||||
printD("sha256: " + hash_value)
|
||||
return hash_value
|
||||
|
||||
|
||||
# get image with full size
|
||||
# width is in number, not string
|
||||
def get_full_size_image_url(image_url, width):
|
||||
return re.sub('/width=\d+/', '/width=' + str(width) + '/', image_url)
|
||||
|
||||
|
||||
# scan model to generate SHA256, then use this SHA256 to get model info from civitai
|
||||
def scan_model(low_memory_sha, max_size_preview, readable_model_info, skip_nsfw_preview):
|
||||
printD("Start scan_model")
|
||||
|
||||
|
||||
count = 0
|
||||
scan_log = ""
|
||||
for model_type, model_folder in model_folders.items():
|
||||
printD("Scanning path: " + model_folder)
|
||||
for root, dirs, files in os.walk(model_folder):
|
||||
for filename in files:
|
||||
# check ext
|
||||
item = os.path.join(root, filename)
|
||||
base, ext = os.path.splitext(item)
|
||||
if ext in model_exts:
|
||||
# find a model
|
||||
# set a Progress log
|
||||
scan_log = "Scanned: " + str(count) + ", Scanning: "+ filename
|
||||
# try to update to UI Here
|
||||
# Is still trying to find a way
|
||||
|
||||
# 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 = ""
|
||||
if low_memory_sha:
|
||||
hash = gen_file_sha256_low_memory(item)
|
||||
else:
|
||||
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 = None
|
||||
if readable_model_info:
|
||||
data = json.dumps(content, indent=4)
|
||||
else:
|
||||
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():
|
||||
img_url = img_dict["url"]
|
||||
if max_size_preview:
|
||||
# use max width
|
||||
if "width" in img_dict.keys():
|
||||
if img_dict["width"]:
|
||||
img_url = get_full_size_image_url(img_url, img_dict["width"])
|
||||
|
||||
printD("Sending request for image: " + img_url)
|
||||
# get image
|
||||
img_r = requests.get(img_url, stream=True)
|
||||
if not img_r.ok:
|
||||
printD("Get error code: " + 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
|
||||
|
||||
# set counter
|
||||
count = count+1
|
||||
|
||||
# for testing, we only check 1 model for each type
|
||||
# break
|
||||
|
||||
scan_log = "Done"
|
||||
|
||||
printD("End scan_model")
|
||||
|
||||
|
||||
|
||||
# handle request from javascript
|
||||
# parameter: msg - msg from js
|
||||
# return: (action, model_type, search_term, 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 "search_term" not in msg_dict.keys():
|
||||
printD("Can not find search_term 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"]
|
||||
search_term = msg_dict["search_term"]
|
||||
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 search_term:
|
||||
printD("search_term 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, search_term, prompt, neg_prompt)
|
||||
|
||||
|
||||
# get model info file's content by model type and model name
|
||||
# parameter: model_type, search_term
|
||||
# return: model_info_dict
|
||||
def get_model_info(model_type, search_term):
|
||||
if model_type not in model_folders.keys():
|
||||
printD("unknow model type: " + model_type)
|
||||
return None
|
||||
|
||||
# search_term = subfolderpath + model name + ext. And it always start with a / even there is no sub folder
|
||||
base, ext = os.path.splitext(search_term)
|
||||
model_info_base = base
|
||||
if base[:1] == "/":
|
||||
model_info_base = base[1:]
|
||||
|
||||
|
||||
model_folder = model_folders[model_type]
|
||||
model_info_filename = model_info_base + civitai_info_suffix + model_info_exts
|
||||
model_info_filepath = os.path.join(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, search_term
|
||||
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, search_term, prompt, neg_prompt = result
|
||||
|
||||
model_info = get_model_info(model_type, search_term)
|
||||
if not model_info:
|
||||
printD(f"Failed to get model info for {model_type} {search_term}")
|
||||
return
|
||||
|
||||
if "modelId" not in model_info.keys():
|
||||
printD(f"Failed to get model id from info file for {model_type} {search_term}")
|
||||
return
|
||||
|
||||
model_id = model_info["modelId"]
|
||||
if not model_id:
|
||||
printD(f"model id from info file of {model_type} {search_term} 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, search_term, 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, search_term, prompt, neg_prompt = result
|
||||
|
||||
|
||||
model_info = get_model_info(model_type, search_term)
|
||||
if not model_info:
|
||||
printD(f"Failed to get model info for {model_type} {search_term}")
|
||||
return [prompt, prompt]
|
||||
|
||||
if "trainedWords" not in model_info.keys():
|
||||
printD(f"Failed to get trainedWords from info file for {model_type} {search_term}")
|
||||
return [prompt, prompt]
|
||||
|
||||
trainedWords = model_info["trainedWords"]
|
||||
if not trainedWords:
|
||||
printD(f"No trainedWords from info file for {model_type} {search_term}")
|
||||
return [prompt, prompt]
|
||||
|
||||
if len(trainedWords) == 0:
|
||||
printD(f"trainedWords from info file for {model_type} {search_term} 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, search_term, prompt, neg_prompt = result
|
||||
|
||||
|
||||
model_info = get_model_info(model_type, search_term)
|
||||
if not model_info:
|
||||
printD(f"Failed to get model info for {model_type} {search_term}")
|
||||
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} {search_term}")
|
||||
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||
|
||||
images = model_info["images"]
|
||||
if not images:
|
||||
printD(f"No images from info file for {model_type} {search_term}")
|
||||
return [prompt, neg_prompt, prompt, neg_prompt]
|
||||
|
||||
if len(images) == 0:
|
||||
printD(f"images from info file for {model_type} {search_term} 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 img["meta"]:
|
||||
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} {search_term} 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:
|
||||
with gr.Blocks() as civitai_helper:
|
||||
# UI will have 3 tabs:
|
||||
# Model Info: Scan model or force a model link to civitai model info by model id or url
|
||||
# Settging: Setting for general use, also can save setting for all tabs
|
||||
# Tool: handy functions, like making all model info readable.
|
||||
with gr.Tab("Model"):
|
||||
with gr.Row():
|
||||
low_memory_sha_ckb = gr.Checkbox(label="Memory Optimised SHA256", value=True, elem_id="ch_low_memory_sha_ckb")
|
||||
max_size_preview_ckb = gr.Checkbox(label="Download Max Size Preview", value=True, elem_id="ch_max_size_preview_ckb")
|
||||
readable_model_info_ckb = gr.Checkbox(label="Readable Model Info file", value=True, elem_id="ch_readable_model_info_ckb")
|
||||
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")
|
||||
|
||||
gr.Markdown("Check console log window for detail, after clicking Scan button")
|
||||
|
||||
|
||||
# with gr.Tab("Settging"):
|
||||
|
||||
# with gr.Tab("Tool"):
|
||||
|
||||
# hidden component for js, not in any tab
|
||||
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=[low_memory_sha_ckb, max_size_preview_ckb, readable_model_info_ckb, 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)
|
||||
|
||||
|
||||
|
|
@ -6,7 +6,7 @@ import re
|
|||
import requests
|
||||
from . import util
|
||||
from . import model
|
||||
|
||||
from . import setting
|
||||
|
||||
suffix = ".civitai"
|
||||
|
||||
|
|
@ -63,6 +63,124 @@ def get_model_info_by_hash(hash:str):
|
|||
|
||||
|
||||
|
||||
def get_model_info_by_id(id:str) -> dict:
|
||||
util.printD("Request model info from civitai")
|
||||
|
||||
if not id:
|
||||
util.printD("id is empty")
|
||||
return
|
||||
|
||||
r = requests.get(url_dict["modelId"]+str(id))
|
||||
if not r.ok:
|
||||
if r.status_code == 404:
|
||||
# this is not a civitai model
|
||||
util.printD("Civitai does not have this model")
|
||||
return {}
|
||||
else:
|
||||
util.printD("Get error code: " + str(r.status_code))
|
||||
util.printD(r.text)
|
||||
return
|
||||
|
||||
# try to get content
|
||||
content = None
|
||||
try:
|
||||
content = r.json()
|
||||
except Exception as e:
|
||||
util.printD("Parse response json failed")
|
||||
util.printD(str(e))
|
||||
util.printD("response:")
|
||||
util.printD(r.text)
|
||||
return
|
||||
|
||||
if not content:
|
||||
util.printD("error, content from civitai is None")
|
||||
return
|
||||
|
||||
return content
|
||||
|
||||
|
||||
def get_version_info_by_version_id(id:str) -> dict:
|
||||
util.printD("Request version info from civitai")
|
||||
|
||||
if not id:
|
||||
util.printD("id is empty")
|
||||
return
|
||||
|
||||
r = requests.get(url_dict["modelVersionId"]+id)
|
||||
if not r.ok:
|
||||
if r.status_code == 404:
|
||||
# this is not a civitai model
|
||||
util.printD("Civitai does not have this model version")
|
||||
return {}
|
||||
else:
|
||||
util.printD("Get error code: " + str(r.status_code))
|
||||
util.printD(r.text)
|
||||
return
|
||||
|
||||
# try to get content
|
||||
content = None
|
||||
try:
|
||||
content = r.json()
|
||||
except Exception as e:
|
||||
util.printD("Parse response json failed")
|
||||
util.printD(str(e))
|
||||
util.printD("response:")
|
||||
util.printD(r.text)
|
||||
return
|
||||
|
||||
if not content:
|
||||
util.printD("error, content from civitai is None")
|
||||
return
|
||||
|
||||
return content
|
||||
|
||||
|
||||
def get_version_info_by_model_id(id:str) -> dict:
|
||||
|
||||
model_info = get_model_info_by_id(id)
|
||||
if not model_info:
|
||||
util.printD(f"Failed to get model info by id: {id}")
|
||||
return
|
||||
|
||||
# check content to get version id
|
||||
if "modelVersions" not in model_info.keys():
|
||||
util.printD("There is no modelVersions in this model_info")
|
||||
return
|
||||
|
||||
if not model_info["modelVersions"]:
|
||||
util.printD("modelVersions is None")
|
||||
return
|
||||
|
||||
if len(model_info["modelVersions"])==0:
|
||||
util.printD("modelVersions is Empty")
|
||||
return
|
||||
|
||||
def_version = model_info["modelVersions"][0]
|
||||
if not def_version:
|
||||
util.printD("default version is None")
|
||||
return
|
||||
|
||||
if "id" not in def_version.keys():
|
||||
util.printD("default version has no id")
|
||||
return
|
||||
|
||||
version_id = def_version["id"]
|
||||
|
||||
if not version_id:
|
||||
util.printD("default version's id is None")
|
||||
return
|
||||
|
||||
# get version info
|
||||
version_info = get_version_info_by_version_id(str(version_id))
|
||||
if not version_info:
|
||||
util.printD(f"Failed to get version info by version_id: {version_id}")
|
||||
return
|
||||
|
||||
return version_info
|
||||
|
||||
|
||||
|
||||
|
||||
# get model info file's content by model type and search_term
|
||||
# parameter: model_type, search_term
|
||||
# return: model_info
|
||||
|
|
@ -87,3 +205,147 @@ def load_model_info_by_search_term(model_type, search_term):
|
|||
return
|
||||
|
||||
return model.load_model_info(model_info_filepath)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# get model file names by model type
|
||||
# parameter: model_type - string
|
||||
# parameter: filter - dict, which kind of model you need
|
||||
# return: model name list
|
||||
def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list:
|
||||
|
||||
model_folder = model.folders[model_type]
|
||||
|
||||
# set filter
|
||||
# only get models don't have a civitai info file
|
||||
no_info_only = False
|
||||
empty_info_only = False
|
||||
|
||||
if filter:
|
||||
if "no_info_only" in filter.keys():
|
||||
no_info_only = filter["no_info_only"]
|
||||
if "empty_info_only" in filter.keys():
|
||||
empty_info_only = filter["empty_info_only"]
|
||||
|
||||
|
||||
|
||||
# get information from filter
|
||||
# only get those model names don't have a civitai model info file
|
||||
model_names = []
|
||||
for root, dirs, files in os.walk(model_folder):
|
||||
for filename in files:
|
||||
item = os.path.join(root, filename)
|
||||
# check extension
|
||||
base, ext = os.path.splitext(item)
|
||||
if ext in model.exts:
|
||||
# find a model
|
||||
|
||||
# check filter
|
||||
if no_info_only:
|
||||
# check model info file
|
||||
info_file = base + suffix + model.info_ext
|
||||
if os.path.isfile(info_file):
|
||||
continue
|
||||
|
||||
if empty_info_only:
|
||||
# check model info file
|
||||
info_file = base + suffix + model.info_ext
|
||||
if os.path.isfile(info_file):
|
||||
# load model info
|
||||
model_info = model.load_model_info(info_file)
|
||||
# check content
|
||||
if model_info:
|
||||
if "id" in model_info.keys():
|
||||
# find a non-empty model info file
|
||||
continue
|
||||
|
||||
model_names.append(filename)
|
||||
|
||||
|
||||
return model_names
|
||||
|
||||
def get_model_names_by_input(model_type, empty_info_only):
|
||||
return get_model_names_by_type_and_filter(model_type, {"empty_info_only":empty_info_only})
|
||||
|
||||
|
||||
# get id from url
|
||||
def get_model_id_from_url(url:str) -> str:
|
||||
util.printD("Run get_model_id_from_url")
|
||||
id = ""
|
||||
|
||||
if not url:
|
||||
util.printD("url or model id can not be empty")
|
||||
return ""
|
||||
|
||||
if url.isnumeric():
|
||||
# is already an id
|
||||
id = url
|
||||
return ""
|
||||
|
||||
s = url.split("/")
|
||||
if len(s) < 2:
|
||||
util.printD("url is not valid")
|
||||
return ""
|
||||
|
||||
if s[-2].isnumeric():
|
||||
id = s[-2]
|
||||
elif s[-1].isnumeric():
|
||||
id = s[-1]
|
||||
else:
|
||||
util.printD("There is no model id in this url")
|
||||
return ""
|
||||
|
||||
return id
|
||||
|
||||
|
||||
# get preview image by model path
|
||||
# image will be saved to file, so no return
|
||||
def get_preview_image_by_model_path(model_path:str, max_size_preview, skip_nsfw_preview) -> str:
|
||||
if not model_path:
|
||||
util.printD("model_path is empty")
|
||||
return
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
util.printD("model_path is not a file")
|
||||
return
|
||||
|
||||
base, ext = os.path.splitext(model_path)
|
||||
first_preview = base+".png"
|
||||
sec_preview = base+".preview.png"
|
||||
info_file = base + suffix + model.info_ext
|
||||
|
||||
# check preview image
|
||||
if not os.path.isfile(sec_preview):
|
||||
# need to download preview image
|
||||
util.printD("Checking preview image for model: " + model_path)
|
||||
# load model_info file
|
||||
if os.path.isfile(info_file):
|
||||
model_info = model.load_model_info(info_file)
|
||||
if not model_info:
|
||||
util.printD("Model Info is empty")
|
||||
return
|
||||
|
||||
if "images" in model_info.keys():
|
||||
if model_info["images"]:
|
||||
for img_dict in model_info["images"]:
|
||||
if "nsfw" in img_dict.keys():
|
||||
if img_dict["nsfw"]:
|
||||
util.printD("This image is NSFW")
|
||||
if skip_nsfw_preview:
|
||||
util.printD("Skip NSFW image")
|
||||
continue
|
||||
|
||||
if "url" in img_dict.keys():
|
||||
img_url = img_dict["url"]
|
||||
if max_size_preview:
|
||||
# use max width
|
||||
if "width" in img_dict.keys():
|
||||
if img_dict["width"]:
|
||||
img_url = get_full_size_image_url(img_url, img_dict["width"])
|
||||
|
||||
util.download_file(img_url, sec_preview)
|
||||
# we only need 1 preview image
|
||||
break
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,11 @@
|
|||
# -*- coding: UTF-8 -*-
|
||||
# handle msg between js and python side
|
||||
import os
|
||||
import json
|
||||
import requests
|
||||
import shutil
|
||||
import webbrowser
|
||||
from . import util
|
||||
from . import model
|
||||
from . import civitai
|
||||
from . import msg_handler
|
||||
|
|
@ -3,110 +3,23 @@
|
|||
import os
|
||||
import json
|
||||
import requests
|
||||
import shutil
|
||||
import webbrowser
|
||||
from . import util
|
||||
from . import model
|
||||
from . import civitai
|
||||
from . import msg
|
||||
from . import msg_handler
|
||||
|
||||
|
||||
|
||||
# scan model to generate SHA256, then use this SHA256 to get model info from civitai
|
||||
def scan_model(low_memory_sha, max_size_preview, skip_nsfw_preview):
|
||||
util.printD("Start scan_model")
|
||||
|
||||
model_count = 0
|
||||
image_count = 0
|
||||
scan_log = ""
|
||||
for model_type, model_folder in model.folders.items():
|
||||
util.printD("Scanning path: " + model_folder)
|
||||
for root, dirs, files in os.walk(model_folder):
|
||||
for filename in files:
|
||||
# check ext
|
||||
item = os.path.join(root, filename)
|
||||
base, ext = os.path.splitext(item)
|
||||
if ext in model.exts:
|
||||
# find a model
|
||||
# set a Progress log
|
||||
scan_log = "Scanned: " + str(model_count) + ", Scanning: "+ filename
|
||||
# try to update to UI Here
|
||||
# Is still trying to find a way
|
||||
|
||||
# get preview image
|
||||
first_preview = base+".png"
|
||||
sec_preview = base+".preview.png"
|
||||
# get info file
|
||||
info_file = base + civitai.suffix + model.info_ext
|
||||
# check info file
|
||||
if not os.path.isfile(info_file):
|
||||
# get model's sha256
|
||||
util.printD("Generate SHA256 for model: " + filename)
|
||||
hash = util.gen_file_sha256(item, low_memory_sha)
|
||||
|
||||
if not hash:
|
||||
util.printD("failed generate SHA256 for this file.")
|
||||
return
|
||||
|
||||
# use this sha256 to get model info from civitai
|
||||
model_info = civitai.get_model_info_by_hash(hash)
|
||||
if model_info is None:
|
||||
util.printD("Fail to get model_info")
|
||||
return
|
||||
|
||||
# write model info to file
|
||||
model.write_model_info(info_file, model_info)
|
||||
|
||||
# set model_count
|
||||
model_count = model_count+1
|
||||
|
||||
# check preview image
|
||||
if not os.path.isfile(sec_preview):
|
||||
# need to download preview image
|
||||
util.printD("Need preview image for this model")
|
||||
# load model_info file
|
||||
if os.path.isfile(info_file):
|
||||
model_info = model.load_model_info(info_file)
|
||||
if not model_info:
|
||||
util.printD("Model Info is empty")
|
||||
continue
|
||||
|
||||
if "images" in model_info.keys():
|
||||
if model_info["images"]:
|
||||
for img_dict in model_info["images"]:
|
||||
if "nsfw" in img_dict.keys():
|
||||
if img_dict["nsfw"]:
|
||||
util.printD("This image is NSFW")
|
||||
if skip_nsfw_preview:
|
||||
util.printD("Skip NSFW image")
|
||||
continue
|
||||
|
||||
if "url" in img_dict.keys():
|
||||
img_url = img_dict["url"]
|
||||
if max_size_preview:
|
||||
# use max width
|
||||
if "width" in img_dict.keys():
|
||||
if img_dict["width"]:
|
||||
img_url = civitai.get_full_size_image_url(img_url, img_dict["width"])
|
||||
|
||||
util.download_file(img_url, sec_preview)
|
||||
image_count = image_count + 1
|
||||
# we only need 1 preview image
|
||||
break
|
||||
|
||||
|
||||
scan_log = "Done"
|
||||
|
||||
util.printD("End scan_model")
|
||||
|
||||
|
||||
|
||||
# get civitai's model url and open it in browser
|
||||
# parameter: model_type, search_term
|
||||
def open_model_url(msg):
|
||||
# output: python msg - will be sent to hidden textbox then picked by js side
|
||||
def open_model_url(msg, open_url_with_js):
|
||||
util.printD("Start open_model_url")
|
||||
|
||||
result = msg.parse_js_msg(msg)
|
||||
output = ""
|
||||
result = msg_handler.parse_js_msg(msg)
|
||||
if not result:
|
||||
util.printD("Parsing js ms failed")
|
||||
return
|
||||
|
|
@ -116,24 +29,36 @@ def open_model_url(msg):
|
|||
model_info = civitai.load_model_info_by_search_term(model_type, search_term)
|
||||
if not model_info:
|
||||
util.printD(f"Failed to get model info for {model_type} {search_term}")
|
||||
return
|
||||
return ""
|
||||
|
||||
if "modelId" not in model_info.keys():
|
||||
util.printD(f"Failed to get model id from info file for {model_type} {search_term}")
|
||||
return
|
||||
return ""
|
||||
|
||||
model_id = model_info["modelId"]
|
||||
if not model_id:
|
||||
util.printD(f"model id from info file of {model_type} {search_term} is None")
|
||||
return
|
||||
return ""
|
||||
|
||||
url = civitai.url_dict["modelPage"]+str(model_id)
|
||||
|
||||
util.printD("Open Url: " + url)
|
||||
# open url
|
||||
webbrowser.open_new_tab(url)
|
||||
|
||||
# msg content for js
|
||||
content = {
|
||||
"url":""
|
||||
}
|
||||
|
||||
if not open_url_with_js:
|
||||
util.printD("Open Url: " + url)
|
||||
# open url
|
||||
webbrowser.open_new_tab(url)
|
||||
else:
|
||||
util.printD("Send Url to js")
|
||||
content["url"] = url
|
||||
output = msg_handler.build_py_msg("open_url", content)
|
||||
|
||||
util.printD("End open_model_url")
|
||||
return output
|
||||
|
||||
|
||||
|
||||
|
|
@ -143,7 +68,7 @@ def open_model_url(msg):
|
|||
def add_trigger_words(msg):
|
||||
util.printD("Start add_trigger_words")
|
||||
|
||||
result = msg.parse_js_msg(msg)
|
||||
result = msg_handler.parse_js_msg(msg)
|
||||
if not result:
|
||||
util.printD("Parsing js ms failed")
|
||||
return
|
||||
|
|
@ -192,7 +117,7 @@ def add_trigger_words(msg):
|
|||
def use_preview_image_prompt(msg):
|
||||
util.printD("Start use_preview_image_prompt")
|
||||
|
||||
result = msg.parse_js_msg(msg)
|
||||
result = msg_handler.parse_js_msg(msg)
|
||||
if not result:
|
||||
util.printD("Parsing js ms failed")
|
||||
return
|
||||
|
|
@ -26,6 +26,8 @@ info_ext = ".info"
|
|||
|
||||
# get cusomter model path
|
||||
def get_custom_model_folder():
|
||||
util.printD("Get Custom Model Folder")
|
||||
|
||||
global folders
|
||||
|
||||
if shared.cmd_opts.embeddings_dir and os.path.isdir(shared.cmd_opts.embeddings_dir):
|
||||
|
|
@ -40,7 +42,7 @@ def get_custom_model_folder():
|
|||
if shared.cmd_opts.lora_dir and os.path.isdir(shared.cmd_opts.lora_dir):
|
||||
folders["lora"] = shared.cmd_opts.lora_dir
|
||||
|
||||
get_custom_model_folder()
|
||||
|
||||
|
||||
|
||||
|
||||
|
|
@ -64,3 +66,55 @@ def load_model_info(path):
|
|||
|
||||
return model_info
|
||||
|
||||
|
||||
# get model file names by model type
|
||||
# parameter: model_type - string
|
||||
# return: model name list
|
||||
def get_model_names_by_type(model_type:str) -> list:
|
||||
|
||||
model_folder = folders[model_type]
|
||||
|
||||
# get information from filter
|
||||
# only get those model names don't have a civitai model info file
|
||||
model_names = []
|
||||
for root, dirs, files in os.walk(model_folder):
|
||||
for filename in files:
|
||||
item = os.path.join(root, filename)
|
||||
# check extension
|
||||
base, ext = os.path.splitext(item)
|
||||
if ext in exts:
|
||||
# find a model
|
||||
model_names.append(filename)
|
||||
|
||||
|
||||
return model_names
|
||||
|
||||
|
||||
# return 2 values: (model_root, model_path)
|
||||
def get_model_path_by_type_and_name(model_type:str, model_name:str) -> str:
|
||||
util.printD("Run get_model_path_by_type_and_name")
|
||||
if model_type not in folders.keys():
|
||||
util.printD("unknown model_type: " + model_type)
|
||||
return
|
||||
|
||||
if not model_name:
|
||||
util.printD("model name can not be empty")
|
||||
return
|
||||
|
||||
folder = folders[model_type]
|
||||
|
||||
# model could be in subfolder, need to walk.
|
||||
model_root = ""
|
||||
model_path = ""
|
||||
for root, dirs, files in os.walk(folder):
|
||||
for filename in files:
|
||||
if filename == model_name:
|
||||
# find model
|
||||
model_root = root
|
||||
model_path = os.path.join(root, filename)
|
||||
return (model_root, model_path)
|
||||
|
||||
return
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,115 @@
|
|||
# -*- coding: UTF-8 -*-
|
||||
# handle msg between js and python side
|
||||
import os
|
||||
from . import util
|
||||
from . import model
|
||||
from . import civitai
|
||||
|
||||
|
||||
|
||||
# scan model to generate SHA256, then use this SHA256 to get model info from civitai
|
||||
# return output msg
|
||||
def scan_model(low_memory_sha, max_size_preview, skip_nsfw_preview):
|
||||
util.printD("Start scan_model")
|
||||
|
||||
output = ""
|
||||
model_count = 0
|
||||
image_count = 0
|
||||
# scan_log = ""
|
||||
for model_type, model_folder in model.folders.items():
|
||||
util.printD("Scanning path: " + model_folder)
|
||||
for root, dirs, files in os.walk(model_folder):
|
||||
for filename in files:
|
||||
# check ext
|
||||
item = os.path.join(root, filename)
|
||||
base, ext = os.path.splitext(item)
|
||||
if ext in model.exts:
|
||||
# find a model
|
||||
# get info file
|
||||
info_file = base + civitai.suffix + model.info_ext
|
||||
# check info file
|
||||
if not os.path.isfile(info_file):
|
||||
# get model's sha256
|
||||
hash = util.gen_file_sha256(item, low_memory_sha)
|
||||
|
||||
if not hash:
|
||||
output = "failed generating SHA256 for model:" + filename
|
||||
util.printD(output)
|
||||
return output
|
||||
|
||||
# use this sha256 to get model info from civitai
|
||||
model_info = civitai.get_model_info_by_hash(hash)
|
||||
if model_info is None:
|
||||
output = "Failed to get model_info"
|
||||
util.printD(output)
|
||||
return output+", check console log for detail"
|
||||
|
||||
# write model info to file
|
||||
model.write_model_info(info_file, model_info)
|
||||
|
||||
# set model_count
|
||||
model_count = model_count+1
|
||||
|
||||
# check preview image
|
||||
civitai.get_preview_image_by_model_path(item, max_size_preview, skip_nsfw_preview)
|
||||
image_count = image_count+1
|
||||
|
||||
|
||||
# scan_log = "Done"
|
||||
|
||||
output = f"Done. Scanned {model_count} models, checked {image_count} images"
|
||||
|
||||
util.printD(output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
|
||||
# Get model info by model type, name and url
|
||||
# output is log info to display on markdown component
|
||||
def get_model_info_by_id(model_type, model_name, model_url_or_id, max_size_preview, skip_nsfw_preview):
|
||||
output = ""
|
||||
# parse model id
|
||||
model_id = civitai.get_model_id_from_url(model_url_or_id)
|
||||
if not model_id:
|
||||
output = "failed to parse model id from url: " + model_url_or_id
|
||||
util.printD(output)
|
||||
return output
|
||||
|
||||
# get model file path
|
||||
# model could be in subfolder
|
||||
result = model.get_model_path_by_type_and_name(model_type, model_name)
|
||||
if not result:
|
||||
output = "failed to get model file path"
|
||||
util.printD(output)
|
||||
return output
|
||||
|
||||
model_root, model_path = result
|
||||
if not model_path:
|
||||
output = "model path is empty"
|
||||
util.printD(output)
|
||||
return output
|
||||
|
||||
# get info file path
|
||||
base, ext = os.path.splitext(model_path)
|
||||
info_file = base + civitai.suffix + model.info_ext
|
||||
|
||||
# get model info
|
||||
#we call it model_info, but in civitai, it is actually version info
|
||||
model_info = civitai.get_version_info_by_model_id(model_id)
|
||||
|
||||
if not model_info:
|
||||
output = "failed to get model info from url: " + model_url_or_id
|
||||
util.printD(output)
|
||||
return output
|
||||
|
||||
# write model info to file
|
||||
model.write_model_info(info_file, model_info)
|
||||
|
||||
util.printD("Saved model info to: "+ info_file)
|
||||
|
||||
# check preview image
|
||||
civitai.get_preview_image_by_model_path(model_path, max_size_preview, skip_nsfw_preview)
|
||||
|
||||
output = "Model Info saved to: " + info_file
|
||||
return output
|
||||
|
|
@ -16,7 +16,7 @@ data = {
|
|||
"skip_nsfw_preview": False
|
||||
},
|
||||
"general":{
|
||||
"open_url_with_js": False,
|
||||
"open_url_with_js": True,
|
||||
"check_model_version_at_startup": False,
|
||||
},
|
||||
"tool":{
|
||||
|
|
@ -26,15 +26,13 @@ data = {
|
|||
|
||||
|
||||
# save setting
|
||||
# return output msg for log
|
||||
def save():
|
||||
print("Saving tranlation service setting...")
|
||||
# write data into globel trans_setting
|
||||
global trans_setting
|
||||
print("Saving setting to: " + path)
|
||||
|
||||
json_data = json.dumps(data, indent=4)
|
||||
|
||||
|
||||
# to json
|
||||
json_data = json.dumps(data)
|
||||
output = ""
|
||||
|
||||
#write to file
|
||||
try:
|
||||
|
|
@ -42,10 +40,14 @@ def save():
|
|||
f.write(json_data)
|
||||
except Exception as e:
|
||||
util.printD("Error when writing file:"+path)
|
||||
output = str(e)
|
||||
util.printD(str(e))
|
||||
return
|
||||
return output
|
||||
|
||||
util.printD("Setting saved to: " + path)
|
||||
output = "Setting saved to: " + path
|
||||
util.printD(output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# load setting to global data
|
||||
|
|
@ -89,7 +91,12 @@ def save_from_input(low_memory_sha, max_size_preview, skip_nsfw_preview, open_ur
|
|||
}
|
||||
}
|
||||
|
||||
save()
|
||||
output = save()
|
||||
|
||||
if not output:
|
||||
output = ""
|
||||
|
||||
return output
|
||||
|
||||
# load to output
|
||||
def load_to_output():
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ def gen_file_sha256(filname, is_low_memory=True):
|
|||
hash_sha256 = hashlib.sha256()
|
||||
with open(filname, "rb") as f:
|
||||
if is_low_memory:
|
||||
printD("Using Memory Optimised SHA256")
|
||||
printD("Using Memory Optimized SHA256")
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_sha256.update(chunk)
|
||||
else:
|
||||
|
|
@ -27,7 +27,7 @@ def gen_file_sha256(filname, is_low_memory=True):
|
|||
|
||||
# get preview image
|
||||
def download_file(url, path):
|
||||
printD("Download file from: " + url)
|
||||
printD("Downloading file from: " + url)
|
||||
# get file
|
||||
r = requests.get(url, stream=True)
|
||||
if not r.ok:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,138 @@
|
|||
# -*- 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 re
|
||||
import modules
|
||||
from modules import script_callbacks
|
||||
from modules import shared
|
||||
from scripts.lib import model
|
||||
from scripts.lib import js_action_civitai
|
||||
from scripts.lib import model_action_civitai
|
||||
from scripts.lib import setting
|
||||
from scripts.lib import civitai
|
||||
|
||||
# init
|
||||
model.get_custom_model_folder()
|
||||
setting.load()
|
||||
|
||||
|
||||
|
||||
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]
|
||||
|
||||
# ====Event's function====
|
||||
def get_model_names_by_input(model_type, empty_info_only):
|
||||
names = civitai.get_model_names_by_input(model_type, empty_info_only)
|
||||
return model_name_drop.update(choices=names)
|
||||
|
||||
|
||||
|
||||
# ====UI====
|
||||
# with gr.Blocks(analytics_enabled=False) as civitai_helper:
|
||||
with gr.Blocks(css="button {background-color: #228be6}") as civitai_helper:
|
||||
|
||||
# init
|
||||
low_memory_sha = setting.data["model"]["low_memory_sha"]
|
||||
max_size_preview = setting.data["model"]["max_size_preview"]
|
||||
skip_nsfw_preview = setting.data["model"]["skip_nsfw_preview"]
|
||||
open_url_with_js = setting.data["general"]["open_url_with_js"]
|
||||
check_model_version_at_startup = setting.data["general"]["check_model_version_at_startup"]
|
||||
|
||||
model_types = list(model.folders.keys())
|
||||
no_info_model_names = civitai.get_model_names_by_input("ckp", False)
|
||||
|
||||
|
||||
# UI will have 3 tabs:
|
||||
# Model Info: Scan model or force a model link to civitai model info by model id or url
|
||||
# General: Setting for general use, also can save setting for all tabs
|
||||
# Tool: handy functions, like making all model info readable.
|
||||
with gr.Tab("Model"):
|
||||
with gr.Box():
|
||||
with gr.Column():
|
||||
gr.Markdown("### Scan Models for Civitai")
|
||||
with gr.Row():
|
||||
low_memory_sha_ckb = gr.Checkbox(label="Memory Optimized SHA256", value=low_memory_sha, elem_id="ch_low_memory_sha_ckb")
|
||||
max_size_preview_ckb = gr.Checkbox(label="Download Max Size Preview", value=max_size_preview, elem_id="ch_max_size_preview_ckb")
|
||||
skip_nsfw_preview_ckb = gr.Checkbox(label="SKip NSFW Preview images", value=skip_nsfw_preview, elem_id="ch_skip_nsfw_preview_ckb")
|
||||
|
||||
# with gr.Row():
|
||||
scan_model_civitai_btn = gr.Button(value="Scan", variant="primary", elem_id="ch_scan_model_civitai_btn")
|
||||
# with gr.Row():
|
||||
scan_model_log_md = gr.Markdown(value="Scanning takes time, just wait. Check console log for detail", elem_id="ch_scan_model_log_md")
|
||||
|
||||
|
||||
with gr.Box():
|
||||
with gr.Column():
|
||||
gr.Markdown("### Get Civitai Model Info by Model ID")
|
||||
with gr.Row():
|
||||
model_type_drop = gr.Dropdown(choices=model_types, label="Model Type", value="ckp", multiselect=False)
|
||||
empty_info_only_ckb = gr.Checkbox(label="Only Show Models have no Info file", value=False, elem_id="cn_empty_info_only_ckb")
|
||||
model_name_drop = gr.Dropdown(choices=no_info_model_names, label="Model", value="ckp", multiselect=False)
|
||||
|
||||
model_url_or_id = gr.Textbox(label="Civitai URL or Model ID", lines=1, value="")
|
||||
get_civitai_model_info_by_id_btn = gr.Button(value="Get 1 Model Info from Civitai")
|
||||
get_model_by_id_log_md = gr.Markdown("")
|
||||
|
||||
|
||||
with gr.Tab("General"):
|
||||
with gr.Row():
|
||||
open_url_with_js_ckb = gr.Checkbox(label="Open Url At Client Side", value=open_url_with_js, elem_id="ch_open_url_with_js_ckb")
|
||||
check_model_version_at_startup_ckb = gr.Checkbox(label="Check Model Version At Startup", value=open_url_with_js, visible=False, elem_id="ch_check_model_version_at_startup_ckb")
|
||||
|
||||
save_setting_btn = gr.Button(value="Save Setting", variant="primary", elem_id="ch_save_setting_btn")
|
||||
general_log_md = gr.Markdown(value="", elem_id="ch_general_log_md")
|
||||
|
||||
# with gr.Tab("Tool"):
|
||||
|
||||
# hidden component for js, not in any tab
|
||||
js_msg_txtbox = gr.Textbox(label="Request Msg From Js", visible=False, lines=1, value="", elem_id="ch_js_msg_txtbox")
|
||||
py_msg_txtbox = gr.Textbox(label="Response Msg From Python", visible=False, lines=1, value="", elem_id="ch_py_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====
|
||||
# Model
|
||||
scan_model_civitai_btn.click(model_action_civitai.scan_model, inputs=[low_memory_sha_ckb, max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=scan_model_log_md)
|
||||
|
||||
model_type_drop.change(get_model_names_by_input, inputs=[model_type_drop, empty_info_only_ckb], outputs=model_name_drop)
|
||||
empty_info_only_ckb.change(get_model_names_by_input, inputs=[model_type_drop, empty_info_only_ckb], outputs=model_name_drop)
|
||||
|
||||
get_civitai_model_info_by_id_btn.click(model_action_civitai.get_model_info_by_id, inputs=[model_type_drop, model_name_drop, model_url_or_id, max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=get_model_by_id_log_md)
|
||||
|
||||
|
||||
# General
|
||||
save_setting_btn.click(setting.save_from_input, inputs=[low_memory_sha_ckb, max_size_preview_ckb, skip_nsfw_preview_ckb, open_url_with_js_ckb, check_model_version_at_startup_ckb], outputs=general_log_md)
|
||||
|
||||
# js action
|
||||
js_open_url_btn.click(js_action_civitai.open_model_url, inputs=[js_msg_txtbox, open_url_with_js_ckb], outputs=py_msg_txtbox)
|
||||
js_add_trigger_words_btn.click(js_action_civitai.add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, img2img_prompt])
|
||||
js_use_preview_prompt_btn.click(js_action_civitai.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