Make custom VAE an Image Setting, rather than a System Setting; Don't load a VAE into memory by default

This commit is contained in:
cmdr2 2022-11-08 16:54:15 +05:30
parent 67cca3bc00
commit 9bc7521de0
8 changed files with 102 additions and 99 deletions

View File

@ -191,7 +191,9 @@ call WHERE uvicorn > .tmp
if not exist "..\models\stable-diffusion" mkdir "..\models\stable-diffusion"
if not exist "..\models\vae" mkdir "..\models\vae"
echo. > "..\models\stable-diffusion\Put your custom ckpt files here.txt"
echo. > "..\models\vae\Put your VAE files here.txt"
@if exist "sd-v1-4.ckpt" (
for %%I in ("sd-v1-4.ckpt") do if "%%~zI" EQU "4265380512" (
@ -321,22 +323,22 @@ echo. > "..\models\stable-diffusion\Put your custom ckpt files here.txt"
@if exist "..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt" (
for %%I in ("..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt") do if "%%~zI" EQU "334695179" (
@if exist "..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt" (
for %%I in ("..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt") do if "%%~zI" EQU "334695179" (
echo "Data files (weights) necessary for the default VAE (sd-vae-ft-mse-original) were already downloaded"
) else (
echo. & echo "The default VAE (sd-vae-ft-mse-original) file present at models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo.
del "..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt"
echo. & echo "The default VAE (sd-vae-ft-mse-original) file present at models\vae\vae-ft-mse-840000-ema-pruned.ckpt is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo.
del "..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt"
)
)
@if not exist "..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt" (
@if not exist "..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt" (
@echo. & echo "Downloading data files (weights) for the default VAE (sd-vae-ft-mse-original).." & echo.
@call curl -L -k https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt > ..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt
@call curl -L -k https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt > ..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt
@if exist "..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt" (
for %%I in ("..\models\stable-diffusion\vae-ft-mse-840000-ema-pruned.vae.pt") do if "%%~zI" NEQ "334695179" (
@if exist "..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt" (
for %%I in ("..\models\vae\vae-ft-mse-840000-ema-pruned.ckpt") do if "%%~zI" NEQ "334695179" (
echo. & echo "Error: The downloaded default VAE (sd-vae-ft-mse-original) file was invalid! Bytes downloaded: %%~zI" & echo.
echo. & echo "Error downloading the data files (weights) for the default VAE (sd-vae-ft-mse-original). Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo.
pause

View File

@ -170,7 +170,9 @@ fi
mkdir -p "../models/stable-diffusion"
mkdir -p "../models/vae"
echo "" > "../models/stable-diffusion/Put your custom ckpt files here.txt"
echo "" > "../models/vae/Put your VAE files here.txt"
if [ -f "sd-v1-4.ckpt" ]; then
model_size=`find "sd-v1-4.ckpt" -printf "%s"`
@ -300,24 +302,24 @@ if [ ! -f "RealESRGAN_x4plus_anime_6B.pth" ]; then
fi
if [ -f "../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt" ]; then
model_size=`find ../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt -printf "%s"`
if [ -f "../models/vae/vae-ft-mse-840000-ema-pruned.ckpt" ]; then
model_size=`find ../models/vae/vae-ft-mse-840000-ema-pruned.ckpt -printf "%s"`
if [ "$model_size" -eq "334695179" ]; then
echo "Data files (weights) necessary for the default VAE (sd-vae-ft-mse-original) were already downloaded"
else
printf "\n\nThe model file present at models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt is invalid. It is only $model_size bytes in size. Re-downloading.."
rm ../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt
printf "\n\nThe model file present at models/vae/vae-ft-mse-840000-ema-pruned.ckpt is invalid. It is only $model_size bytes in size. Re-downloading.."
rm ../models/vae/vae-ft-mse-840000-ema-pruned.ckpt
fi
fi
if [ ! -f "../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt" ]; then
if [ ! -f "../models/vae/vae-ft-mse-840000-ema-pruned.ckpt" ]; then
echo "Downloading data files (weights) for the default VAE (sd-vae-ft-mse-original).."
curl -L -k https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt > ../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt
curl -L -k https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt > ../models/vae/vae-ft-mse-840000-ema-pruned.ckpt
if [ -f "../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt" ]; then
model_size=`find ../models/stable-diffusion/vae-ft-mse-840000-ema-pruned.vae.pt -printf "%s"`
if [ -f "../models/vae/vae-ft-mse-840000-ema-pruned.ckpt" ]; then
model_size=`find ../models/vae/vae-ft-mse-840000-ema-pruned.ckpt -printf "%s"`
if [ ! "$model_size" -eq "334695179" ]; then
printf "\n\nError: The downloaded default VAE (sd-vae-ft-mse-original) file was invalid! Bytes downloaded: $model_size\n\n"
printf "\n\nError downloading the data files (weights) for the default VAE (sd-vae-ft-mse-original). Sorry about that, please try to:\n 1. Run this installer again.\n 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting\n 3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB\n 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues\nThanks!\n\n"

View File

@ -18,7 +18,7 @@
<div id="container">
<div id="top-nav">
<div id="logo">
<h1>Stable Diffusion UI <small>v2.3.9 <span id="updateBranchLabel"></span></small></h1>
<h1>Stable Diffusion UI <small>v2.3.10 <span id="updateBranchLabel"></span></small></h1>
</div>
<ul id="top-nav-items">
<li class="dropdown">
@ -100,6 +100,11 @@
<!-- <option value="sd-v1-4" selected>sd-v1-4</option> -->
</select>
</td></tr>
<tr class="pl-5"><td><label for="vae_model">Custom VAE:</label></td><td>
<select id="vae_model" name="vae_model">
<!-- <option value="" selected>None</option> -->
</select>
</td></tr>
<tr id="samplerSelection" class="pl-5"><td><label for="sampler">Sampler:</label></td><td>
<select id="sampler" name="sampler">
<option value="plms">plms</option>
@ -255,13 +260,13 @@
</div>
</body>
<script src="media/js/parameters.js?v=1"></script>
<script src="media/js/parameters.js?v=2"></script>
<script src="media/js/plugins.js?v=1"></script>
<script src="media/js/utils.js?v=5"></script>
<script src="media/js/inpainting-editor.js?v=1"></script>
<script src="media/js/image-modifiers.js?v=4"></script>
<script src="media/js/auto-save.js?v=3"></script>
<script src="media/js/main.js?v=7"></script>
<script src="media/js/auto-save.js?v=4"></script>
<script src="media/js/main.js?v=8"></script>
<script src="media/js/themes.js?v=3"></script>
<script src="media/js/dnd.js?v=3"></script>
<script>

View File

@ -13,6 +13,7 @@ const SETTINGS_IDS_LIST = [
"num_outputs_total",
"num_outputs_parallel",
"stable_diffusion_model",
"vae_model",
"sampler",
"width",
"height",

View File

@ -39,7 +39,7 @@ let useFaceCorrectionField = document.querySelector("#use_face_correction")
let useUpscalingField = document.querySelector("#use_upscale")
let upscaleModelField = document.querySelector("#upscale_model")
let stableDiffusionModelField = document.querySelector('#stable_diffusion_model')
let vaeModelField = document.querySelector('#default_vae_model')
let vaeModelField = document.querySelector('#vae_model')
let outputFormatField = document.querySelector('#output_format')
let showOnlyFilteredImageField = document.querySelector("#show_only_filtered_image")
let updateBranchLabel = document.querySelector("#updateBranchLabel")
@ -774,6 +774,7 @@ function getCurrentUserRequest() {
use_cpu: useCPUField.checked,
use_full_precision: useFullPrecisionField.checked,
use_stable_diffusion_model: stableDiffusionModelField.value,
use_vae_model: vaeModelField.value,
stream_progress_updates: true,
stream_image_progress: (numOutputsTotal > 50 ? false : streamImageProgressField.checked),
show_only_filtered_image: showOnlyFilteredImageField.checked,
@ -823,7 +824,10 @@ function makeImage() {
function createTask(task) {
let taskConfig = `Seed: ${task.seed}, Sampler: ${task.reqBody.sampler}, Inference Steps: ${task.reqBody.num_inference_steps}, Guidance Scale: ${task.reqBody.guidance_scale}, Model: ${task.reqBody.use_stable_diffusion_model}`
if (negativePromptField.value.trim() !== '') {
if (task.reqBody.use_vae_model.trim() !== '') {
taskConfig += `, VAE: ${task.reqBody.use_vae_model}`
}
if (task.reqBody.negative_prompt.trim() !== '') {
taskConfig += `, Negative Prompt: ${task.reqBody.negative_prompt}`
}
if (task.reqBody.init_image !== undefined) {
@ -1139,12 +1143,6 @@ useBetaChannelField.addEventListener('click', async function(e) {
})
})
vaeModelField.addEventListener('change', async function() {
await changeAppConfig({
'model_vae': this.value
})
})
async function getAppConfig() {
try {
let res = await fetch('/get/app_config')
@ -1163,22 +1161,25 @@ async function getAppConfig() {
async function getModels() {
try {
var model_setting_key = "stable_diffusion_model"
var selectedModel = SETTINGS[model_setting_key].value
var sd_model_setting_key = "stable_diffusion_model"
var vae_model_setting_key = "vae_model"
var selectedSDModel = SETTINGS[sd_model_setting_key].value
var selectedVaeModel = SETTINGS[vae_model_setting_key].value
let res = await fetch('/get/models')
const models = await res.json()
let activeModels = models['active']
console.log('get models response', models)
let modelOptions = models['options']
let stableDiffusionOptions = modelOptions['stable-diffusion']
let vaeOptions = modelOptions['vae']
let activeVae = activeModels['vae']
vaeOptions.unshift('') // add a None option
function createModelOptions(modelField, selectedModel) {
return function(modelName) {
let modelOption = document.createElement('option')
modelOption.value = modelName
modelOption.innerText = modelName
modelOption.innerText = modelName !== '' ? modelName : 'None'
if (modelName === selectedModel) {
modelOption.selected = true
@ -1188,16 +1189,14 @@ async function getModels() {
}
}
stableDiffusionOptions.forEach(createModelOptions(stableDiffusionModelField, selectedModel))
vaeOptions.forEach(createModelOptions(vaeModelField, activeVae))
stableDiffusionOptions.forEach(createModelOptions(stableDiffusionModelField, selectedSDModel))
vaeOptions.forEach(createModelOptions(vaeModelField, selectedVaeModel))
// TODO: set default for model here too
SETTINGS[model_setting_key].default = stableDiffusionOptions[0]
if (getSetting(model_setting_key) == '' || SETTINGS[model_setting_key].value == '') {
setSetting(model_setting_key, stableDiffusionOptions[0])
SETTINGS[sd_model_setting_key].default = stableDiffusionOptions[0]
if (getSetting(sd_model_setting_key) == '' || SETTINGS[sd_model_setting_key].value == '') {
setSetting(sd_model_setting_key, stableDiffusionOptions[0])
}
console.log('get models response', models)
} catch (e) {
console.log('get models error', e)
}

View File

@ -51,11 +51,6 @@ var PARAMETERS = [
return `<input id="${parameter.id}" name="${parameter.id}" size="40" disabled>`
}
},
{
id: "default_vae_model",
type: ParameterType.select, // Note: options generated dynamically
label: "Default VAE",
},
{
id: "sound_toggle",
type: ParameterType.checkbox,

View File

@ -258,13 +258,15 @@ def load_model_ckpt():
_, _ = modelFS.load_state_dict(sd, strict=False)
if thread_data.vae_file is not None:
if os.path.exists(thread_data.vae_file + '.vae.pt'):
print(f"Loading VAE weights from: {thread_data.vae_file}.vae.pt")
vae_ckpt = torch.load(thread_data.vae_file + '.vae.pt', map_location="cpu")
vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss"}
modelFS.first_stage_model.load_state_dict(vae_dict, strict=False)
else:
print(f'Cannot find VAE file: {thread_data.vae_file}.vae.pt')
for model_extension in ['.ckpt', '.vae.pt']:
if os.path.exists(thread_data.vae_file + model_extension):
print(f"Loading VAE weights from: {thread_data.vae_file}{model_extension}")
vae_ckpt = torch.load(thread_data.vae_file + model_extension, map_location="cpu")
vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss"}
modelFS.first_stage_model.load_state_dict(vae_dict, strict=False)
break
else:
print(f'Cannot find VAE file: {thread_data.vae_file}{model_extension}')
modelFS.eval()
if thread_data.device != 'cpu':

View File

@ -30,9 +30,6 @@ APP_CONFIG_DEFAULT_MODELS = [
'custom-model', # Check if user has a custom model, use it first.
'sd-v1-4', # Default fallback.
]
APP_CONFIG_DEFAULT_VAE = [
'vae-ft-mse-840000-ema-pruned', # Default fallback.
]
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
@ -142,7 +139,7 @@ def setConfig(config):
except Exception as e:
print(traceback.format_exc())
def resolve_model_to_use(model_name:str, model_type:str, model_dir:str, model_extension:str, default_models=[]):
def resolve_model_to_use(model_name:str, model_type:str, model_dir:str, model_extensions:list, default_models=[]):
model_dirs = [os.path.join(MODELS_DIR, model_dir), SD_DIR]
if not model_name: # When None try user configured model.
config = getConfig()
@ -151,43 +148,38 @@ def resolve_model_to_use(model_name:str, model_type:str, model_dir:str, model_ex
if model_name:
# Check models directory
models_dir_path = os.path.join(MODELS_DIR, model_dir, model_name)
if os.path.exists(models_dir_path + model_extension):
return models_dir_path
if os.path.exists(model_name + model_extension):
# Direct Path to file
model_name = os.path.abspath(model_name)
return model_name
for model_extension in model_extensions:
if os.path.exists(models_dir_path + model_extension):
return models_dir_path
if os.path.exists(model_name + model_extension):
# Direct Path to file
model_name = os.path.abspath(model_name)
return model_name
# Default locations
if model_name in default_models:
default_model_path = os.path.join(SD_DIR, model_name)
if os.path.exists(default_model_path + model_extension):
return default_model_path
for model_extension in model_extensions:
if os.path.exists(default_model_path + model_extension):
return default_model_path
# Can't find requested model, check the default paths.
for default_model in default_models:
for model_dir in model_dirs:
default_model_path = os.path.join(model_dir, default_model)
if os.path.exists(default_model_path + model_extension):
if model_name is not None:
print(f'Could not find the configured custom model {model_name}{model_extension}. Using the default one: {default_model_path}{model_extension}')
return default_model_path
for model_extension in model_extensions:
if os.path.exists(default_model_path + model_extension):
if model_name is not None:
print(f'Could not find the configured custom model {model_name}{model_extension}. Using the default one: {default_model_path}{model_extension}')
return default_model_path
raise Exception('No valid models found.')
def resolve_ckpt_to_use(model_name:str=None):
return resolve_model_to_use(model_name, model_type='stable-diffusion', model_dir='stable-diffusion', model_extension='.ckpt', default_models=APP_CONFIG_DEFAULT_MODELS)
return resolve_model_to_use(model_name, model_type='stable-diffusion', model_dir='stable-diffusion', model_extensions=['.ckpt'], default_models=APP_CONFIG_DEFAULT_MODELS)
def resolve_vae_to_use(ckpt_model_path:str=None):
if ckpt_model_path is not None:
if os.path.exists(ckpt_model_path + '.vae.pt'):
return ckpt_model_path
ckpt_model_name = os.path.basename(ckpt_model_path)
model_dirs = [os.path.join(MODELS_DIR, 'stable-diffusion'), SD_DIR]
for model_dir in model_dirs:
default_model_path = os.path.join(model_dir, ckpt_model_name)
if os.path.exists(default_model_path + '.vae.pt'):
return default_model_path
return resolve_model_to_use(model_name=None, model_type='vae', model_dir='stable-diffusion', model_extension='.vae.pt', default_models=APP_CONFIG_DEFAULT_VAE)
def resolve_vae_to_use(model_name:str=None):
try:
return resolve_model_to_use(model_name, model_type='vae', model_dir='vae', model_extensions=['.vae.pt', '.ckpt'], default_models=[])
except:
return None
class SetAppConfigRequest(BaseModel):
update_branch: str = None
@ -213,10 +205,6 @@ async def setAppConfig(req : SetAppConfigRequest):
render_devices.append('GPU:' + req.render_devices)
if len(render_devices) > 0:
config['render_devices'] = render_devices
if req.model_vae:
if 'model' not in config:
config['model'] = {}
config['model']['vae'] = req.model_vae
try:
setConfig(config)
return JSONResponse({'status': 'OK'}, headers=NOCACHE_HEADERS)
@ -236,21 +224,25 @@ def getModels():
},
}
def listModels(models_dirname, model_type, model_extensions):
models_dir = os.path.join(MODELS_DIR, models_dirname)
for file in os.listdir(models_dir):
for model_extension in model_extensions:
if file.endswith(model_extension):
model_name = file[:-len(model_extension)]
models['options'][model_type].append(model_name)
models['options'][model_type] = [*set(models['options'][model_type])] # remove duplicates
# custom models
sd_models_dir = os.path.join(MODELS_DIR, 'stable-diffusion')
for model_type, model_extension in [('stable-diffusion', '.ckpt'), ('vae', '.vae.pt')]:
for file in os.listdir(sd_models_dir):
if file.endswith(model_extension):
model_name = file[:-len(model_extension)]
models['options'][model_type].append(model_name)
listModels(models_dirname='stable-diffusion', model_type='stable-diffusion', model_extensions=['.ckpt'])
listModels(models_dirname='vae', model_type='vae', model_extensions=['.vae.pt', '.ckpt'])
# legacy
custom_weight_path = os.path.join(SD_DIR, 'custom-model.ckpt')
if os.path.exists(custom_weight_path):
models['options']['stable-diffusion'].append('custom-model')
models['active']['vae'] = os.path.basename(task_manager.default_vae_to_load)
return models
def getUIPlugins():
@ -307,12 +299,17 @@ def ping(session_id:str=None):
response['session'] = 'pending'
return JSONResponse(response, headers=NOCACHE_HEADERS)
def save_model_to_config(model_name):
def save_model_to_config(ckpt_model_name, vae_model_name):
config = getConfig()
if 'model' not in config:
config['model'] = {}
config['model']['stable-diffusion'] = model_name
config['model']['stable-diffusion'] = ckpt_model_name
config['model']['vae'] = vae_model_name
if vae_model_name is None or vae_model_name == "":
del config['model']['vae']
setConfig(config)
@app.post('/render')
@ -325,9 +322,9 @@ def render(req : task_manager.ImageRequest):
if req.use_face_correction and task_manager.is_alive(0) <= 0: #TODO Remove when GFPGANer is fixed upstream.
raise HTTPException(status_code=412, detail=f'GFPGANer only works GPU:0, use CUDA_VISIBLE_DEVICES if GFPGANer is needed on a specific GPU.') # HTTP412 Precondition Failed
try:
save_model_to_config(req.use_stable_diffusion_model)
save_model_to_config(req.use_stable_diffusion_model, req.use_vae_model)
req.use_stable_diffusion_model = resolve_ckpt_to_use(req.use_stable_diffusion_model)
req.use_vae_model = resolve_vae_to_use(ckpt_model_path=req.use_stable_diffusion_model)
req.use_vae_model = resolve_vae_to_use(req.use_vae_model)
new_task = task_manager.render(req)
response = {
'status': str(task_manager.current_state),
@ -406,7 +403,7 @@ config = getConfig()
# Start the task_manager
task_manager.default_model_to_load = resolve_ckpt_to_use()
task_manager.default_vae_to_load = resolve_vae_to_use(ckpt_model_path=task_manager.default_model_to_load)
task_manager.default_vae_to_load = resolve_vae_to_use()
if 'render_devices' in config: # Start a new thread for each device.
if isinstance(config['render_devices'], str):
config['render_devices'] = config['render_devices'].split(',')