mirror of
https://github.com/sd-webui/stable-diffusion-webui.git
synced 2024-12-14 06:35:14 +03:00
120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
import streamlit as st
|
|
|
|
import warnings
|
|
import os
|
|
import k_diffusion as K
|
|
from omegaconf import OmegaConf
|
|
|
|
from sd_utils import *
|
|
|
|
#streamlit components section
|
|
from st_on_hover_tabs import on_hover_tabs
|
|
|
|
|
|
try:
|
|
# this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
|
|
from transformers import logging
|
|
|
|
logging.set_verbosity_error()
|
|
except:
|
|
pass
|
|
|
|
# remove some annoying deprecation warnings that show every now and then.
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
|
|
st.session_state["defaults"] = OmegaConf.load("configs/webui/webui_streamlit.yaml")
|
|
if (os.path.exists("configs/webui/userconfig_streamlit.yaml")):
|
|
user_defaults = OmegaConf.load("configs/webui/userconfig_streamlit.yaml");
|
|
st.session_state["defaults"] = OmegaConf.merge(st.session_state["defaults"], user_defaults)
|
|
|
|
# this should force GFPGAN and RealESRGAN onto the selected gpu as well
|
|
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" # see issue #152
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(st.session_state["defaults"].general.gpu)
|
|
|
|
# functions to load css locally OR remotely starts here. Options exist for future flexibility. Called as st.markdown with unsafe_allow_html as css injection
|
|
# TODO, maybe look into async loading the file especially for remote fetching
|
|
def local_css(file_name):
|
|
with open(file_name) as f:
|
|
st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
|
|
|
|
def remote_css(url):
|
|
st.markdown(f'<link href="{url}" rel="stylesheet">', unsafe_allow_html=True)
|
|
|
|
def load_css(isLocal, nameOrURL):
|
|
if(isLocal):
|
|
local_css(nameOrURL)
|
|
else:
|
|
remote_css(nameOrURL)
|
|
|
|
def layout():
|
|
"""Layout functions to define all the streamlit layout here."""
|
|
st.set_page_config(page_title="Stable Diffusion Playground", layout="wide")
|
|
|
|
with st.empty():
|
|
# load css as an external file, function has an option to local or remote url. Potential use when running from cloud infra that might not have access to local path.
|
|
load_css(True, 'frontend/css/streamlit.main.css')
|
|
|
|
# check if the models exist on their respective folders
|
|
if os.path.exists(os.path.join(st.session_state["defaults"].general.GFPGAN_dir, "experiments", "pretrained_models", "GFPGANv1.3.pth")):
|
|
st.session_state["GFPGAN_available"] = True
|
|
else:
|
|
st.session_state["GFPGAN_available"] = False
|
|
|
|
if os.path.exists(os.path.join(st.session_state["defaults"].general.RealESRGAN_dir, "experiments","pretrained_models", f"{st.session_state['defaults'].general.RealESRGAN_model}.pth")):
|
|
st.session_state["RealESRGAN_available"] = True
|
|
else:
|
|
st.session_state["RealESRGAN_available"] = False
|
|
|
|
# Allow for custom models to be used instead of the default one,
|
|
# an example would be Waifu-Diffusion or any other fine tune of stable diffusion
|
|
st.session_state["custom_models"]:sorted = []
|
|
for root, dirs, files in os.walk(os.path.join("models", "custom")):
|
|
for file in files:
|
|
if os.path.splitext(file)[1] == '.ckpt':
|
|
#fullpath = os.path.join(root, file)
|
|
#print(fullpath)
|
|
st.session_state["custom_models"].append(os.path.splitext(file)[0])
|
|
#print (os.path.splitext(file)[0])
|
|
|
|
if len(st.session_state["custom_models"]) > 0:
|
|
st.session_state["CustomModel_available"] = True
|
|
st.session_state["custom_models"].append("Stable Diffusion v1.4")
|
|
else:
|
|
st.session_state["CustomModel_available"] = False
|
|
|
|
with st.sidebar:
|
|
# The global settings section will be moved to the Settings page.
|
|
#with st.expander("Global Settings:"):
|
|
#st.write("Global Settings:")
|
|
#defaults.general.update_preview = st.checkbox("Update Image Preview", value=defaults.general.update_preview,
|
|
#help="If enabled the image preview will be updated during the generation instead of at the end. You can use the Update Preview \
|
|
#Frequency option bellow to customize how frequent it's updated. By default this is enabled and the frequency is set to 1 step.")
|
|
#st.session_state.update_preview_frequency = st.text_input("Update Image Preview Frequency", value=defaults.general.update_preview_frequency,
|
|
#help="Frequency in steps at which the the preview image is updated. By default the frequency is set to 1 step.")
|
|
|
|
tabs = on_hover_tabs(tabName=['Stable Diffusion', "Textual Inversion","Model Manager","Settings"],
|
|
iconName=['dashboard','model_training' ,'cloud_download', 'settings'], default_choice=0)
|
|
|
|
if tabs =='Stable Diffusion':
|
|
txt2img_tab, img2img_tab, txt2vid_tab, postprocessing_tab = st.tabs(["Text-to-Image Unified", "Image-to-Image Unified",
|
|
"Text-to-Video","Post-Processing"])
|
|
with txt2img_tab:
|
|
from txt2img import layout
|
|
layout()
|
|
|
|
with img2img_tab:
|
|
from img2img import layout
|
|
layout()
|
|
|
|
with txt2vid_tab:
|
|
from txt2vid import layout
|
|
layout()
|
|
|
|
|
|
#
|
|
elif tabs == 'Model Manager':
|
|
from ModelManager import layout
|
|
layout()
|
|
|
|
if __name__ == '__main__':
|
|
layout() |