173 lines
5.4 KiB
Python
173 lines
5.4 KiB
Python
# -*- coding: UTF-8 -*-
|
|
# handle msg between js and python side
|
|
import os
|
|
import json
|
|
import requests
|
|
import webbrowser
|
|
from . import util
|
|
from . import model
|
|
from . import civitai
|
|
from . import msg_handler
|
|
|
|
|
|
|
|
|
|
# get civitai's model url and open it in browser
|
|
# parameter: model_type, search_term
|
|
# 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")
|
|
|
|
output = ""
|
|
result = msg_handler.parse_js_msg(msg)
|
|
if not result:
|
|
util.printD("Parsing js ms failed")
|
|
return
|
|
|
|
action, model_type, search_term, prompt, neg_prompt = result
|
|
|
|
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 ""
|
|
|
|
if "modelId" not in model_info.keys():
|
|
util.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:
|
|
util.printD(f"model id from info file of {model_type} {search_term} is None")
|
|
return ""
|
|
|
|
url = civitai.url_dict["modelPage"]+str(model_id)
|
|
|
|
|
|
# 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
|
|
|
|
|
|
|
|
# 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):
|
|
util.printD("Start add_trigger_words")
|
|
|
|
result = msg_handler.parse_js_msg(msg)
|
|
if not result:
|
|
util.printD("Parsing js ms failed")
|
|
return
|
|
|
|
action, model_type, search_term, prompt, neg_prompt = result
|
|
|
|
|
|
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 [prompt, prompt]
|
|
|
|
if "trainedWords" not in model_info.keys():
|
|
util.printD(f"Failed to get trainedWords from info file for {model_type} {search_term}")
|
|
return [prompt, prompt]
|
|
|
|
trainedWords = model_info["trainedWords"]
|
|
if not trainedWords:
|
|
util.printD(f"No trainedWords from info file for {model_type} {search_term}")
|
|
return [prompt, prompt]
|
|
|
|
if len(trainedWords) == 0:
|
|
util.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
|
|
util.printD("trigger_words: " + trigger_words)
|
|
util.printD("prompt: " + prompt)
|
|
util.printD("new_prompt: " + new_prompt)
|
|
|
|
util.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):
|
|
util.printD("Start use_preview_image_prompt")
|
|
|
|
result = msg_handler.parse_js_msg(msg)
|
|
if not result:
|
|
util.printD("Parsing js ms failed")
|
|
return
|
|
|
|
action, model_type, search_term, prompt, neg_prompt = result
|
|
|
|
|
|
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 [prompt, neg_prompt, prompt, neg_prompt]
|
|
|
|
if "images" not in model_info.keys():
|
|
util.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:
|
|
util.printD(f"No images from info file for {model_type} {search_term}")
|
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
|
|
|
if len(images) == 0:
|
|
util.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:
|
|
util.printD(f"There is no prompt of {model_type} {search_term} in its preview image")
|
|
return [prompt, neg_prompt, prompt, neg_prompt]
|
|
|
|
util.printD("End use_preview_image_prompt")
|
|
|
|
return [preview_prompt, preview_neg_prompt, preview_prompt, preview_neg_prompt]
|
|
|
|
|