mirror of
https://github.com/easydiffusion/easydiffusion.git
synced 2025-02-16 10:29:39 +01:00
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:
parent
67cca3bc00
commit
9bc7521de0
@ -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
|
||||
|
@ -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"
|
||||
|
@ -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>
|
||||
|
@ -13,6 +13,7 @@ const SETTINGS_IDS_LIST = [
|
||||
"num_outputs_total",
|
||||
"num_outputs_parallel",
|
||||
"stable_diffusion_model",
|
||||
"vae_model",
|
||||
"sampler",
|
||||
"width",
|
||||
"height",
|
||||
|
@ -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)
|
||||
}
|
||||
|
@ -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,
|
||||
|
@ -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':
|
||||
|
91
ui/server.py
91
ui/server.py
@ -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(',')
|
||||
|
Loading…
Reference in New Issue
Block a user