mirror of
https://github.com/sd-webui/stable-diffusion-webui.git
synced 2024-12-14 14:52:31 +03:00
257 lines
9.3 KiB
Python
257 lines
9.3 KiB
Python
# This file is part of stable-diffusion-webui (https://github.com/sd-webui/stable-diffusion-webui/).
|
|
|
|
# Copyright 2022 sd-webui team.
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
# base webui import and utils.
|
|
from sd_utils import *
|
|
|
|
# streamlit imports
|
|
import streamlit.components.v1 as components
|
|
|
|
|
|
class plugin_info():
|
|
plugname = "concept_library"
|
|
description = "Concept Library"
|
|
displayPriority = 4
|
|
|
|
|
|
# Init Vuejs component
|
|
_component_func = components.declare_component(
|
|
"sd-concepts-browser", "./frontend/dists/concept-browser/dist")
|
|
|
|
|
|
def sdConceptsBrowser(concepts, key=None):
|
|
component_value = _component_func(concepts=concepts, key=key, default="")
|
|
return component_value
|
|
|
|
|
|
@st.experimental_memo(persist="disk", show_spinner=False, suppress_st_warning=True)
|
|
def getConceptsFromPath(page, conceptPerPage, searchText=""):
|
|
#print("getConceptsFromPath", "page:", page, "conceptPerPage:", conceptPerPage, "searchText:", searchText)
|
|
# get the path where the concepts are stored
|
|
path = os.path.join(
|
|
os.getcwd(), st.session_state['defaults'].general.sd_concepts_library_folder)
|
|
acceptedExtensions = ('jpeg', 'jpg', "png")
|
|
concepts = []
|
|
|
|
if os.path.exists(path):
|
|
# List all folders (concepts) in the path
|
|
folders = [f for f in os.listdir(
|
|
path) if os.path.isdir(os.path.join(path, f))]
|
|
filteredFolders = folders
|
|
|
|
# Filter the folders by the search text
|
|
if searchText != "":
|
|
filteredFolders = [
|
|
f for f in folders if searchText.lower() in f.lower()]
|
|
else:
|
|
filteredFolders = []
|
|
|
|
conceptIndex = 1
|
|
for folder in filteredFolders:
|
|
# handle pagination
|
|
if conceptIndex > (page * conceptPerPage):
|
|
continue
|
|
if conceptIndex <= ((page - 1) * conceptPerPage):
|
|
conceptIndex += 1
|
|
continue
|
|
|
|
concept = {
|
|
"name": folder,
|
|
"token": "<" + folder + ">",
|
|
"images": [],
|
|
"type": ""
|
|
}
|
|
|
|
# type of concept is inside type_of_concept.txt
|
|
typePath = os.path.join(path, folder, "type_of_concept.txt")
|
|
binFile = os.path.join(path, folder, "learned_embeds.bin")
|
|
|
|
# Continue if the concept is not valid or the download has failed (no type_of_concept.txt or no binFile)
|
|
if not os.path.exists(typePath) or not os.path.exists(binFile):
|
|
continue
|
|
|
|
with open(typePath, "r") as f:
|
|
concept["type"] = f.read()
|
|
|
|
# List all files in the concept/concept_images folder
|
|
files = [f for f in os.listdir(os.path.join(path, folder, "concept_images")) if os.path.isfile(
|
|
os.path.join(path, folder, "concept_images", f))]
|
|
# Retrieve only the 4 first images
|
|
for file in files:
|
|
|
|
# Skip if we already have 4 images
|
|
if len(concept["images"]) >= 4:
|
|
break
|
|
|
|
if file.endswith(acceptedExtensions):
|
|
try:
|
|
# Add a copy of the image to avoid file locking
|
|
originalImage = Image.open(os.path.join(
|
|
path, folder, "concept_images", file))
|
|
|
|
# Maintain the aspect ratio (max 200x200)
|
|
resizedImage = originalImage.copy()
|
|
resizedImage.thumbnail((200, 200), Image.ANTIALIAS)
|
|
|
|
# concept["images"].append(resizedImage)
|
|
|
|
concept["images"].append(imageToBase64(resizedImage))
|
|
# Close original image
|
|
originalImage.close()
|
|
except:
|
|
print("Error while loading image", file, "in concept", folder, "(The file may be corrupted). Skipping it.")
|
|
|
|
concepts.append(concept)
|
|
conceptIndex += 1
|
|
# print all concepts name
|
|
#print("Results:", [c["name"] for c in concepts])
|
|
return concepts
|
|
|
|
@st.cache(persist=True, allow_output_mutation=True, show_spinner=False, suppress_st_warning=True)
|
|
def imageToBase64(image):
|
|
import io
|
|
import base64
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
return img_str
|
|
|
|
|
|
@st.experimental_memo(persist="disk", show_spinner=False, suppress_st_warning=True)
|
|
def getTotalNumberOfConcepts(searchText=""):
|
|
# get the path where the concepts are stored
|
|
path = os.path.join(
|
|
os.getcwd(), st.session_state['defaults'].general.sd_concepts_library_folder)
|
|
concepts = []
|
|
|
|
if os.path.exists(path):
|
|
# List all folders (concepts) in the path
|
|
folders = [f for f in os.listdir(
|
|
path) if os.path.isdir(os.path.join(path, f))]
|
|
filteredFolders = folders
|
|
|
|
# Filter the folders by the search text
|
|
if searchText != "":
|
|
filteredFolders = [
|
|
f for f in folders if searchText.lower() in f.lower()]
|
|
else:
|
|
filteredFolders = []
|
|
return len(filteredFolders)
|
|
|
|
|
|
|
|
def layout():
|
|
# 2 tabs, one for Concept Library and one for the Download Manager
|
|
tab_library, tab_downloader = st.tabs(["Library", "Download Manager"])
|
|
|
|
# Concept Library
|
|
with tab_library:
|
|
downloaded_concepts_count = getTotalNumberOfConcepts()
|
|
concepts_per_page = st.session_state["defaults"].concepts_library.concepts_per_page
|
|
|
|
if not "results" in st.session_state:
|
|
st.session_state["results"] = getConceptsFromPath(1, concepts_per_page, "")
|
|
|
|
# Pagination controls
|
|
if not "cl_current_page" in st.session_state:
|
|
st.session_state["cl_current_page"] = 1
|
|
|
|
# Search
|
|
if not 'cl_search_text' in st.session_state:
|
|
st.session_state["cl_search_text"] = ""
|
|
|
|
if not 'cl_search_results_count' in st.session_state:
|
|
st.session_state["cl_search_results_count"] = downloaded_concepts_count
|
|
|
|
# Search bar
|
|
_search_col, _refresh_col = st.columns([10, 2])
|
|
with _search_col:
|
|
search_text_input = st.text_input("Search", "", placeholder=f'Search for a concept ({downloaded_concepts_count} available)', label_visibility="hidden")
|
|
if search_text_input != st.session_state["cl_search_text"]:
|
|
# Search text has changed
|
|
st.session_state["cl_search_text"] = search_text_input
|
|
st.session_state["cl_current_page"] = 1
|
|
st.session_state["cl_search_results_count"] = getTotalNumberOfConcepts(st.session_state["cl_search_text"])
|
|
st.session_state["results"] = getConceptsFromPath(1, concepts_per_page, st.session_state["cl_search_text"])
|
|
|
|
with _refresh_col:
|
|
# Super weird fix to align the refresh button with the search bar ( Please streamlit, add css support.. )
|
|
_refresh_col.write("")
|
|
_refresh_col.write("")
|
|
if st.button("Refresh concepts", key="refresh_concepts", help="Refresh the concepts folders. Use this if you have added new concepts manually or deleted some."):
|
|
getTotalNumberOfConcepts.clear()
|
|
getConceptsFromPath.clear()
|
|
st.experimental_rerun()
|
|
|
|
|
|
# Show results
|
|
results_empty = st.empty()
|
|
|
|
# Pagination
|
|
pagination_empty = st.empty()
|
|
|
|
# Layouts
|
|
with pagination_empty:
|
|
with st.container():
|
|
if len(st.session_state["results"]) > 0:
|
|
last_page = math.ceil(st.session_state["cl_search_results_count"] / concepts_per_page)
|
|
_1, _2, _3, _4, _previous_page, _current_page, _next_page, _9, _10, _11, _12 = st.columns([1,1,1,1,1,2,1,1,1,1,1])
|
|
|
|
# Previous page
|
|
with _previous_page:
|
|
if st.button("Previous", key="cl_previous_page"):
|
|
st.session_state["cl_current_page"] -= 1
|
|
if st.session_state["cl_current_page"] <= 0:
|
|
st.session_state["cl_current_page"] = last_page
|
|
st.session_state["results"] = getConceptsFromPath(st.session_state["cl_current_page"], concepts_per_page, st.session_state["cl_search_text"])
|
|
|
|
# Current page
|
|
with _current_page:
|
|
_current_page_container = st.empty()
|
|
|
|
# Next page
|
|
with _next_page:
|
|
if st.button("Next", key="cl_next_page"):
|
|
st.session_state["cl_current_page"] += 1
|
|
if st.session_state["cl_current_page"] > last_page:
|
|
st.session_state["cl_current_page"] = 1
|
|
st.session_state["results"] = getConceptsFromPath(st.session_state["cl_current_page"], concepts_per_page, st.session_state["cl_search_text"])
|
|
|
|
# Current page
|
|
with _current_page_container:
|
|
st.markdown(f'<p style="text-align: center">Page {st.session_state["cl_current_page"]} of {last_page}</p>', unsafe_allow_html=True)
|
|
# st.write(f"Page {st.session_state['cl_current_page']} of {last_page}", key="cl_current_page")
|
|
|
|
with results_empty:
|
|
with st.container():
|
|
if downloaded_concepts_count == 0:
|
|
st.write("You don't have any concepts in your library ")
|
|
st.markdown("To add concepts to your library, download some from the [sd-concepts-library](https://github.com/sd-webui/sd-concepts-library) \
|
|
repository and save the content of `sd-concepts-library` into ```./models/custom/sd-concepts-library``` or just create your own concepts :wink:.", unsafe_allow_html=False)
|
|
else:
|
|
if len(st.session_state["results"]) == 0:
|
|
st.write("No concept found in the library matching your search: " + st.session_state["cl_search_text"])
|
|
else:
|
|
# display number of results
|
|
if st.session_state["cl_search_text"]:
|
|
st.write(f"Found {st.session_state['cl_search_results_count']} {'concepts' if st.session_state['cl_search_results_count'] > 1 else 'concept' } matching your search")
|
|
sdConceptsBrowser(st.session_state['results'], key="results")
|
|
|
|
|
|
with tab_downloader:
|
|
st.write("Not implemented yet")
|
|
|
|
return False
|