|
|
import ast |
|
|
import enum |
|
|
import os |
|
|
import sys |
|
|
from typing import Any, Optional |
|
|
import gradio as gr |
|
|
import importlib |
|
|
import json |
|
|
import urllib.parse |
|
|
from fastapi import FastAPI |
|
|
from modules import script_callbacks, shared |
|
|
|
|
|
prefix = "/whapi/v1" |
|
|
|
|
|
allowed_preview_extensions = ["png", "jpg", "jpeg", "webp", "gif"] |
|
|
|
|
|
class UI_TYPE(enum.Enum): |
|
|
Unknown = 1 |
|
|
A1111 = 2 |
|
|
NEW_FORGE = 3 |
|
|
|
|
|
class WeightHelperAPI: |
|
|
|
|
|
instance = None |
|
|
|
|
|
network = None |
|
|
network_lora = None |
|
|
ui_type = None |
|
|
|
|
|
@staticmethod |
|
|
def init_endpoints(_: gr.Blocks, app: FastAPI): |
|
|
if not WeightHelperAPI.instance: |
|
|
WeightHelperAPI.instance = WeightHelperAPI() |
|
|
|
|
|
instance = WeightHelperAPI.instance |
|
|
@app.post(prefix + "/get_template") |
|
|
async def _(): |
|
|
return instance.get_template() |
|
|
@app.post(prefix + "/get_metadata") |
|
|
async def _(key: str): |
|
|
instance.__initialize() |
|
|
return instance.get_metadata(key) |
|
|
@app.post(prefix + "/get_preview_info") |
|
|
async def _(key: str): |
|
|
instance.__initialize() |
|
|
return instance.get_preview_info(key) |
|
|
|
|
|
def __new__(cls, *args, **kwargs): |
|
|
if not cls.instance: |
|
|
cls.instance = super(WeightHelperAPI, cls).__new__(cls, *args, **kwargs) |
|
|
return cls.instance |
|
|
|
|
|
def __initialize(self): |
|
|
if self.network_lora is not None or self.ui_type is not None: |
|
|
return |
|
|
|
|
|
try: |
|
|
self.network = importlib.import_module("extensions-builtin.Lora.network") |
|
|
self.network_lora = importlib.import_module("extensions-builtin.Lora.ui_extra_networks_lora").networks |
|
|
self.ui_type = UI_TYPE.A1111 |
|
|
except: |
|
|
try: |
|
|
self.network = importlib.import_module("packages_3rdparty.webui_lora_collection.network") |
|
|
|
|
|
self.network_lora = importlib.import_module("extensions-builtin.sd_forge_lora.networks") |
|
|
self.ui_type = UI_TYPE.NEW_FORGE |
|
|
except Exception as e: |
|
|
self.ui_type = UI_TYPE.Unknown |
|
|
print(e) |
|
|
|
|
|
def __get_lora_on_disk(self, key): |
|
|
if self.network_lora: |
|
|
lora_on_disk = self.network_lora.available_network_aliases.get(key) |
|
|
if not lora_on_disk: |
|
|
|
|
|
lora_on_disk = next((v for v in self.network_lora.available_network_aliases.values() if v.alias == key), None) |
|
|
return lora_on_disk |
|
|
return None |
|
|
|
|
|
def get_template(self): |
|
|
template_path = os.path.join(__file__, "../../html/template.hbs") |
|
|
try: |
|
|
with open(template_path, "r", encoding="utf-8", errors="replace") as f: |
|
|
return f.read() |
|
|
except OSError: |
|
|
pass |
|
|
|
|
|
def get_metadata(self, key) -> dict[str, Optional[str]]: |
|
|
ret: dict[str, Optional[str]] = { |
|
|
"sd_version": None, |
|
|
"algorithm": None, |
|
|
"base_model": None |
|
|
} |
|
|
|
|
|
lora_on_disk = self.__get_lora_on_disk(key) |
|
|
if lora_on_disk: |
|
|
ret["sd_version"] = self.__get_sd_version(lora_on_disk) |
|
|
|
|
|
metadata = lora_on_disk.metadata |
|
|
ss_network_args = self.__get_network_args(lora_on_disk) |
|
|
ret["algorithm"] = self.__get_algorithm(metadata, ss_network_args) |
|
|
|
|
|
model_path, _ = os.path.splitext(lora_on_disk.filename) |
|
|
lora_json = self.__get_lora_json(model_path) |
|
|
json_sd_version = lora_json.get("sd version", None) |
|
|
if json_sd_version: |
|
|
ret["sd_version"] = json_sd_version |
|
|
|
|
|
civitai_info = self.__find_civitai_info(model_path) |
|
|
ret["base_model"] = civitai_info.get("baseModel", None) |
|
|
|
|
|
return ret |
|
|
|
|
|
def get_preview_info(self, key): |
|
|
ret: dict[str, Optional[bool | str | list[str]]] = { |
|
|
"model_id": None, |
|
|
"trigger_words": None, |
|
|
"negative_trigger_words": [], |
|
|
"model_name": None, |
|
|
"preview_url": None, |
|
|
"has_metadata": False, |
|
|
"description": None, |
|
|
} |
|
|
|
|
|
lora_on_disk = self.__get_lora_on_disk(key) |
|
|
if lora_on_disk: |
|
|
model_path, _ = os.path.splitext(lora_on_disk.filename) |
|
|
ret["model_name"] = lora_on_disk.name |
|
|
ret["preview_url"] = self.__find_preview(model_path) |
|
|
ret["has_metadata"] = self.__find_metadata(lora_on_disk) |
|
|
ret["description"] = self.__find_description(model_path) |
|
|
|
|
|
trigger_words = [] |
|
|
civitai_info = self.__find_civitai_info(model_path) |
|
|
ret["model_id"] = civitai_info.get("modelId", None) |
|
|
trigger_words = civitai_info.get("trainedWords", []) |
|
|
trigger_words = ",".join(trigger_words).split(",") |
|
|
ret["trigger_words"] = [w.strip() for w in trigger_words if w] |
|
|
|
|
|
lora_json = self.__get_lora_json(model_path) |
|
|
|
|
|
activation_text = lora_json.get("activation text", "") |
|
|
if activation_text: |
|
|
activation_text = activation_text.split(",") |
|
|
activation_text = [w.strip() for w in activation_text if w] |
|
|
ret["trigger_words"] = activation_text |
|
|
|
|
|
negative_text = lora_json.get("negative text", "") |
|
|
if negative_text: |
|
|
negative_text = negative_text.split(",") |
|
|
negative_text = [w.strip() for w in negative_text if w] |
|
|
ret["negative_trigger_words"] = negative_text |
|
|
|
|
|
json_description = lora_json.get("description", None) |
|
|
if json_description: |
|
|
ret["description"] = json_description |
|
|
|
|
|
return ret |
|
|
|
|
|
def __get_sd_version(self, lora) -> Optional[str]: |
|
|
if self.ui_type == UI_TYPE.A1111: |
|
|
sd_version: Optional[str] = lora.sd_version.name |
|
|
if sd_version and self.network: |
|
|
if sd_version == self.network.SdVersion.SD1.name or sd_version == self.network.SdVersion.SD2.name: |
|
|
return "SD" |
|
|
elif sd_version == self.network.SdVersion.SDXL.name: |
|
|
return "SDXL" |
|
|
return sd_version |
|
|
elif self.ui_type == UI_TYPE.NEW_FORGE: |
|
|
base_model_version: Optional[str] = lora.metadata.get("ss_base_model_version") |
|
|
if base_model_version: |
|
|
base_model_version = base_model_version.lower() |
|
|
if "sd_" in base_model_version: |
|
|
return "SD" |
|
|
if "sd3" in base_model_version: |
|
|
return "SD" |
|
|
elif "sdxl" in base_model_version: |
|
|
return "SDXL" |
|
|
elif "flux" in base_model_version: |
|
|
return "FLUX" |
|
|
return None |
|
|
return None |
|
|
|
|
|
def __get_algorithm(self, metadata, ss_network_args) -> Optional[str]: |
|
|
ss_network_module = metadata.get("ss_network_module") |
|
|
if not ss_network_module or ss_network_module == "Unknown": |
|
|
return None |
|
|
|
|
|
conv_dim = float(ss_network_args.get("conv_dim", "-1")) |
|
|
conv_alpha = float(ss_network_args.get("conv_alpha", "-1")) |
|
|
algo = ss_network_args.get("algo", "").lower() |
|
|
unit = ss_network_args.get("unit", "").lower() |
|
|
dora_wd = ss_network_args.get("dora_wd", False) |
|
|
|
|
|
if ss_network_module.find("locon.locon_kohya") >= 0: |
|
|
return "LoCon" |
|
|
|
|
|
elif ss_network_module.find("lycoris.kohya") >= 0: |
|
|
algoName = "" |
|
|
if dora_wd: |
|
|
algoName = "(DoRA)" |
|
|
elif algo == "lora": |
|
|
algoName += "(LoCon)" |
|
|
elif algo == "locon": |
|
|
if unit: |
|
|
algoName += "(DyLoRA)" |
|
|
else: |
|
|
algoName += "(LoCon)" |
|
|
elif algo == "loha": |
|
|
algoName += "(LoHa)" |
|
|
elif algo == "lokr": |
|
|
algoName += "(Lokr)" |
|
|
elif algo == "ia3": |
|
|
algoName += "(IA3)" |
|
|
elif algo == "full": |
|
|
algoName += "(Full)" |
|
|
elif algo == "glora": |
|
|
algoName += "(GLoRA)" |
|
|
|
|
|
return f"LyCORIS{algoName}" |
|
|
|
|
|
elif ss_network_module.find("networks.dylora") >= 0: |
|
|
if algo == "dylora" and unit: |
|
|
if conv_dim > 0 or conv_alpha > 0: |
|
|
return "DyLoRA(C3Lier)" |
|
|
return "DyLoRA(LierLa)" |
|
|
|
|
|
elif conv_dim > 0 or conv_alpha > 0: |
|
|
return "LoRA(C3Lier)" |
|
|
|
|
|
else: |
|
|
return "LoRA(LierLa)" |
|
|
|
|
|
def __get_network_args(self, lora) -> dict[str, Any]: |
|
|
metadata = lora.metadata |
|
|
ss_network_args = metadata.get("ss_network_args", {}) |
|
|
if type(ss_network_args) is str: |
|
|
try: |
|
|
ss_network_args = ast.literal_eval(ss_network_args) |
|
|
except Exception: |
|
|
ss_network_args = {} |
|
|
return ss_network_args |
|
|
|
|
|
def __get_lora_json(self, path) -> dict[str, Any]: |
|
|
if path: |
|
|
lora_json = f"{path}.json" |
|
|
if os.path.exists(lora_json): |
|
|
try: |
|
|
with open(lora_json, "r", encoding="utf-8", errors="replace") as f: |
|
|
return json.load(f) |
|
|
except OSError: |
|
|
pass |
|
|
return {} |
|
|
|
|
|
def __link_preview(self, filename) -> str: |
|
|
quoted_filename = urllib.parse.quote(filename.replace('\\', '/')) |
|
|
|
|
|
|
|
|
|
|
|
return f"./sd_extra_networks/thumb?filename={quoted_filename}" |
|
|
|
|
|
def __find_preview(self, path) -> str: |
|
|
if path: |
|
|
potential_files = sum([[f"{path}.{ext}", f"{path}.preview.{ext}"] for ext in allowed_preview_extensions], []) |
|
|
for file in potential_files: |
|
|
if os.path.exists(file): |
|
|
return self.__link_preview(file) |
|
|
return "./file=html/card-no-preview.png" |
|
|
|
|
|
def __find_metadata(self, lora) -> bool: |
|
|
return len(lora.metadata) > 0 if lora else False |
|
|
|
|
|
|
|
|
def __find_description(self, path) -> Optional[str]: |
|
|
if path: |
|
|
for file in [f"{path}.txt", f"{path}.description.txt"]: |
|
|
if os.path.exists(file): |
|
|
try: |
|
|
with open(file, "r", encoding="utf-8", errors="replace") as f: |
|
|
return f.read() |
|
|
except OSError: |
|
|
pass |
|
|
return None |
|
|
|
|
|
|
|
|
def __find_civitai_info(self, path) -> dict[str, Any]: |
|
|
if path: |
|
|
civitai_info = f"{path}.civitai.info" |
|
|
if os.path.exists(civitai_info): |
|
|
try: |
|
|
with open(civitai_info, "r", encoding="utf-8", errors="replace") as f: |
|
|
return json.load(f) |
|
|
except OSError: |
|
|
pass |
|
|
return {} |
|
|
|
|
|
def on_ui_settings(): |
|
|
""" |
|
|
Set up the UI settings for the weight helper by adding options to the shared configuration. |
|
|
""" |
|
|
section = ('weight_helper', 'Weight Helper') |
|
|
|
|
|
shared.options_templates.update(shared.options_section(section, { |
|
|
'weight_helper_enabled': shared.OptionInfo(True, 'Enabled'), |
|
|
'weight_helper_context_menu_scale': shared.OptionInfo(0.8, "Context Menu Scale", gr.Number), |
|
|
'weight_helper_using_execCommand': shared.OptionInfo(True, "Use the deprecated execCommand function to replace text." |
|
|
).info( |
|
|
"You can use 'undo' to revert the text to its previous state, " |
|
|
"but it will no longer be updated in real-time." |
|
|
), |
|
|
|
|
|
"weight_helper_slider_length": shared.OptionInfo(160, "Slider Length (px)", gr.Number), |
|
|
|
|
|
"weight_helper_te_min": shared.OptionInfo(0, "TEnc Min Value", gr.Number), |
|
|
"weight_helper_te_max": shared.OptionInfo(1, "TEnc Max Value", gr.Number), |
|
|
"weight_helper_te_step": shared.OptionInfo(0.05, "TEnc Step", gr.Number), |
|
|
|
|
|
"weight_helper_unet_min": shared.OptionInfo(0, "UNet Min Value", gr.Number), |
|
|
"weight_helper_unet_max": shared.OptionInfo(1, "UNet Max Value", gr.Number), |
|
|
"weight_helper_unet_step": shared.OptionInfo(0.05, "UNet Step", gr.Number), |
|
|
|
|
|
"weight_helper_dyn_min": shared.OptionInfo(0, "Dyn Min Value", gr.Number), |
|
|
"weight_helper_dyn_max": shared.OptionInfo(128, "Dyn Max Value", gr.Number), |
|
|
"weight_helper_dyn_step": shared.OptionInfo(8, "Dyn Step", gr.Number), |
|
|
|
|
|
"weight_helper_lbw_min": shared.OptionInfo(0, "LBW (LoRA Block Weight) Min Value", gr.Number), |
|
|
"weight_helper_lbw_max": shared.OptionInfo(1, "LBW (LoRA Block Weight) Max Value", gr.Number), |
|
|
"weight_helper_lbw_step": shared.OptionInfo(0.05, "LBW (LoRA Block Weight) Step", gr.Number), |
|
|
|
|
|
'weight_helper_show_preview': shared.OptionInfo(True, 'Show Preview'), |
|
|
"weight_helper_preview_height": shared.OptionInfo(400, "Preview Height (px)", gr.Number), |
|
|
"weight_helper_preview_position": shared.OptionInfo("Top Right", "Preview Position", gr.Radio, { |
|
|
"choices": ["Top Right", "Bottom Right", "Top Left", "Bottom Left"] |
|
|
}), |
|
|
|
|
|
"weight_helper_lbw_lora_sd_block_points": shared.OptionInfo( |
|
|
"BASE, IN01-IN04, IN05-IN08, M00, OUT03-OUT06, OUT07-OUT11", |
|
|
"Advanced option - LoRA-LierLa block points" |
|
|
).info( |
|
|
"default: BASE, IN01-IN04, IN05-IN08, M00, OUT03-OUT06, OUT07-OUT11" |
|
|
), |
|
|
"weight_helper_lbw_lycoris_sd_block_points": shared.OptionInfo( |
|
|
"BASE, IN00-IN05, IN06-IN11, M00, OUT00-OUT05, OUT06-OUT11", |
|
|
"Advanced option - LyCORIS or LoRA-C3Lier block points" |
|
|
).info( |
|
|
"default: BASE, IN00-IN05, IN06-IN11, M00, OUT00-OUT05, OUT06-OUT11" |
|
|
), |
|
|
"weight_helper_lbw_lora_sdxl_block_points": shared.OptionInfo( |
|
|
"BASE, IN04-IN08, M00, OUT00-OUT05", |
|
|
"Advanced option - SDXL LoRA-LierLa block points" |
|
|
).info( |
|
|
"default: BASE, IN04-IN08, M00, OUT00-OUT05" |
|
|
), |
|
|
"weight_helper_lbw_lycoris_sdxl_block_points": shared.OptionInfo( |
|
|
"BASE, IN00-IN03, IN04-IN08, M00, OUT00-OUT03, OUT04-OUT08", |
|
|
"Advanced option - SDXL LyCORIS or LoRA-C3Lier block points" |
|
|
).info( |
|
|
"default: BASE, IN00-IN03, IN04-IN08, M00, OUT00-OUT03, OUT04-OUT08" |
|
|
) |
|
|
})) |
|
|
|
|
|
script_callbacks.on_ui_settings(on_ui_settings) |
|
|
script_callbacks.on_app_started(WeightHelperAPI.init_endpoints) |
|
|
|
