add load unet override

This commit is contained in:
Vladimir Mandic
2024-04-28 11:51:08 -04:00
parent 01e2ab051f
commit a26d222cc1
10 changed files with 83 additions and 18 deletions
+1 -1
View File
@@ -2,7 +2,7 @@
"parserOptions": {
"ecmaVersion": 2020
},
"plugins": ["html", "json", "markdown"],
"plugins": ["html", "json"],
"extends": [
"plugin:json/recommended",
"eslint:recommended",
+5
View File
@@ -63,6 +63,11 @@ it is time to give credit to original [author](https://github.com/auTOMATIC1111)
> --ipadapter 'Plus:~/generative/Samples/cutie-512.png:0.5'
- Add API endpoint `/sdapi/v1/vqa` and CLI util `cli/simple-vqa.py`
- **Models**:
- **Load UNET**: ability to override/load external UNET to a selected model
Works similar to how VAE is selected and loaded: Set UNet folder and UNet model in settings
Can be replaced on-the-fly, not just during initial model load
Enables usage of fine-tunes such as [DPO-SD15](https://huggingface.co/mhdang/dpo-sd1.5-text2image-v1) or [DPO-SDXL](https://huggingface.co/mhdang/dpo-sdxl-text2image-v1)
*Note*: if there is a `JSON` file with the same name as the model it will be used as Unet config, otherwise Unet config from currently loaded model will be used
- [PixArt-Σ](https://pixart-alpha.github.io/PixArt-sigma-project/)
pixart-Σ is a high end diffusion Transformer model (DiT) with a T5 encoder/decoder capable of directly generating images at 4K resolution
to use, simply select from *networks -> models -> PixArt-Σ*
+6 -6
View File
@@ -54,11 +54,11 @@
--line-md: 1.4em;
--line-lg: 1.5em;
--range-shadow:
-20em 0 0 0 hsl(180, 54%, 21%), -19em 0 0 0 hsl(180, 54%, 23%), -18em 0 0 0 hsl(180, 54%, 25%), -17em 0 0 0 hsl(180, 54%, 27%),
-16em 0 0 0 hsl(180, 54%, 29%), -15em 0 0 0 hsl(180, 54%, 31%), -14em 0 0 0 hsl(180, 54%, 33%), -13em 0 0 0 hsl(180, 54%, 35%),
-12em 0 0 0 hsl(180, 54%, 37%), -11em 0 0 0 hsl(180, 54%, 39%), -10em 0 0 0 hsl(180, 54%, 41%), -9em 0 0 0 hsl(180, 54%, 43%),
-8em 0 0 0 hsl(180, 54%, 45%), -7em 0 0 0 hsl(180, 54%, 47%), -6em 0 0 0 hsl(180, 54%, 49%), -5em 0 0 0 hsl(180, 54%, 51%),
-4em 0 0 0 hsl(180, 54%, 53%), -3em 0 0 0 hsl(180, 54%, 55%), -2em 0 0 0 hsl(180, 54%, 57%), -1em 0 0 0 hsl(180, 54%, 59%),
-20em 0 0 0 hsl(180, 54%, 2%), -19em 0 0 0 hsl(180, 54%, 5%), -18em 0 0 0 hsl(180, 54%, 0%), -17em 0 0 0 hsl(180, 54%, 11%),
-16em 0 0 0 hsl(180, 54%, 14%), -15em 0 0 0 hsl(180, 54%, 17%), -14em 0 0 0 hsl(180, 54%, 20%), -13em 0 0 0 hsl(180, 54%, 23%),
-12em 0 0 0 hsl(180, 54%, 26%), -11em 0 0 0 hsl(180, 54%, 29%), -10em 0 0 0 hsl(180, 54%, 32%), -9em 0 0 0 hsl(180, 54%, 35%),
-8em 0 0 0 hsl(180, 54%, 38%), -7em 0 0 0 hsl(180, 54%, 41%), -6em 0 0 0 hsl(180, 54%, 44%), -5em 0 0 0 hsl(180, 54%, 47%),
-4em 0 0 0 hsl(180, 54%, 50%), -3em 0 0 0 hsl(180, 54%, 53%), -2em 0 0 0 hsl(180, 54%, 56%), -1em 0 0 0 hsl(180, 54%, 59%),
1em 0 0 0 var(--neutral-800), 2em 0 0 0 var(--neutral-800), 3em 0 0 0 var(--neutral-800), 4em 0 0 0 var(--neutral-800),
5em 0 0 0 var(--neutral-800), 6em 0 0 0 var(--neutral-800), 7em 0 0 0 var(--neutral-800), 8em 0 0 0 var(--neutral-800),
9em 0 0 0 var(--neutral-800), 10em 0 0 0 var(--neutral-800), 11em 0 0 0 var(--neutral-800), 12em 0 0 0 var(--neutral-800),
@@ -78,7 +78,7 @@ input[type=range]::-webkit-slider-thumb { border: 0px solid #000000 !important;
input[type=range]::-moz-range-thumb { border: 0px solid #000000 !important; height: var(--line-xs) !important; width: var(--line-md) !important; border-radius: var(--radius-lg) !important; background: var(--highlight-color) !important; cursor: pointer !important; appearance: none !important; margin-top: -4px !important; }
*/
input[type='range'] { display: block; margin: 0; padding: 0; height: 1em; background-color: transparent; overflow: hidden; cursor: pointer; box-shadow: 0 0 0 0 transparent; -webkit-appearance: none; appearance: none; }
input[type='range']::-webkit-slider-thumb { height: .9em; width: .9em; background-color: hsl(180, 54%, 61%); box-shadow: var(--range-shadow); border-radius: var(--radius-sm); }
input[type='range']::-webkit-slider-thumb { height: .9em; width: .9em; background-color: hsl(180, 54%, 61%); box-shadow: var(--range-shadow); border-radius: var(--radius-xs); }
input[type='range']::-webkit-slider-runnable-track, input[type='range']::-webkit-slider-thumb { -webkit-appearance: none; }
input[type='range']::-moz-range-thumb { height: .9em; width: .9em; background-color: hsl(180, 54%, 61%); box-shadow: var(--range-shadow); }
input[type='range']::-moz-range-track, input[type='range']::-webkit-slider-runnable-track { border: none; background: none; width: 100%; height: 100%; }
+1
View File
@@ -100,6 +100,7 @@ def create_paths(opts):
create_path(fix_path('ckpt_dir'))
create_path(fix_path('diffusers_dir'))
create_path(fix_path('vae_dir'))
create_path(fix_path('unet_dir'))
create_path(fix_path('lora_dir'))
create_path(fix_path('embeddings_dir'))
create_path(fix_path('hypernetwork_dir'))
+10 -10
View File
@@ -20,7 +20,7 @@ from omegaconf import OmegaConf
import tomesd
from transformers import logging as transformers_logging
from ldm.util import instantiate_from_config
from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, errors, hashes, sd_models_config, sd_models_compile, sd_hijack_accelerate
from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, hashes, sd_models_config, sd_models_compile, sd_hijack_accelerate
from modules.timer import Timer
from modules.memstats import memory_stats
from modules.modeldata import model_data
@@ -1113,6 +1113,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
if hasattr(sd_model, "set_progress_bar_config"):
sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba')
sd_unet.load_unet(sd_model)
set_diffuser_options(sd_model, vae, op)
if op == 'refiner' and shared.opts.diffusers_move_refiner:
@@ -1134,15 +1135,14 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
shared.log.error("Failed to load diffusers model")
errors.display(e, "loading Diffusers model")
if sd_model is not None:
from modules.textual_inversion import textual_inversion
sd_model.embedding_db = textual_inversion.EmbeddingDatabase()
if op == 'refiner':
model_data.sd_refiner = sd_model
else:
model_data.sd_model = sd_model
sd_model.embedding_db.add_embedding_dir(shared.opts.embeddings_dir)
sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True)
from modules.textual_inversion import textual_inversion
sd_model.embedding_db = textual_inversion.EmbeddingDatabase()
if op == 'refiner':
model_data.sd_refiner = sd_model
else:
model_data.sd_model = sd_model
sd_model.embedding_db.add_embedding_dir(shared.opts.embeddings_dir)
sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True)
timer.record("load")
devices.torch_gc(force=True)
+42
View File
@@ -0,0 +1,42 @@
import os
from modules import shared, devices, files_cache
unet_dict = {}
def load_unet(model):
from diffusers import UNet2DConditionModel
from safetensors.torch import load_file
if shared.opts.sd_unet == 'None':
return
if shared.opts.sd_unet not in list(unet_dict):
shared.log.error(f'UNet model not found: {shared.opts.sd_unet}')
return
if not hasattr(model, 'unet') or model.unet is None:
shared.log.error('UNet not found in current model')
return
config_file = os.path.splitext(unet_dict[shared.opts.sd_unet])[0] + '.json'
if os.path.exists(config_file):
config = shared.readfile(config_file)
else:
config = None
config_file = 'default'
try:
shared.log.info(f'Loading UNet: name="{shared.opts.sd_unet}" file="{unet_dict[shared.opts.sd_unet]}" config="{config_file}"')
unet = UNet2DConditionModel.from_config(model.unet.config if config is None else config).to(devices.device, devices.dtype)
state_dict = load_file(unet_dict[shared.opts.sd_unet])
unet.load_state_dict(state_dict)
model.unet = unet.to(devices.device, devices.dtype_unet)
except Exception as e:
unet = None
shared.log.error(f'Failed to load UNet model: {e}')
return
def refresh_unet_list():
unet_dict.clear()
for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors"]):
name = os.path.splitext(os.path.basename(file))[0]
unet_dict[name] = file
shared.log.debug(f'Available UNets: path="{shared.opts.unet_dir}" items={len(unet_dict)}')
+2
View File
@@ -386,6 +386,7 @@ options_templates.update(options_section(('sd', "Execution & Models"), {
"sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
"sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints),
"sd_vae": OptionInfo("Automatic", "VAE model", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list),
"sd_unet": OptionInfo("None", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list),
"sd_checkpoint_autoload": OptionInfo(True, "Model autoload on start"),
"sd_model_dict": OptionInfo('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints),
"stream_load": OptionInfo(False, "Load models using stream loading method", gr.Checkbox, {"visible": backend == Backend.ORIGINAL }),
@@ -539,6 +540,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), {
"diffusers_dir": OptionInfo(os.path.join(paths.models_path, 'Diffusers'), "Folder with Huggingface models", folder=True),
"hfcache_dir": OptionInfo(os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub'), "Folder for Huggingface cache", folder=True),
"vae_dir": OptionInfo(os.path.join(paths.models_path, 'VAE'), "Folder with VAE files", folder=True),
"unet_dir": OptionInfo(os.path.join(paths.models_path, 'UNET'), "Folder with UNET files", folder=True),
"sd_lora": OptionInfo("", "Add LoRA to prompt", gr.Textbox, {"visible": False}),
"lora_dir": OptionInfo(os.path.join(paths.models_path, 'Lora'), "Folder with LoRA network(s)", folder=True),
"lyco_dir": OptionInfo(os.path.join(paths.models_path, 'LyCORIS'), "Folder with LyCORIS network(s)", gr.Text, {"visible": False}),
+10
View File
@@ -13,6 +13,16 @@ def refresh_vae_list():
modules.sd_vae.refresh_vae_list()
def sd_unet_items():
import modules.sd_unet
return ["None"] + list(modules.sd_unet.unet_dict)
def refresh_unet_list():
import modules.sd_unet
modules.sd_unet.refresh_unet_list()
def list_crossattention(diffusers=False):
if diffusers:
return [
+5
View File
@@ -23,6 +23,7 @@ import modules.lowvram
import modules.scripts
import modules.sd_models
import modules.sd_vae
import modules.sd_unet
import modules.progress
import modules.ui
import modules.txt2img
@@ -93,6 +94,9 @@ def initialize():
modules.sd_vae.refresh_vae_list()
timer.startup.record("vae")
modules.sd_unet.refresh_unet_list()
timer.startup.record("unet")
extensions.list_extensions()
timer.startup.record("extensions")
@@ -117,6 +121,7 @@ def initialize():
timer.startup.record("upscalers")
shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False)
shared.opts.onchange("sd_unet", wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False)
shared.opts.onchange("temp_dir", gr_tempdir.on_tmpdir_changed)
timer.startup.record("onchange")