Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

optimized model downloading

Browse Source
This commit is contained in:
Seunghoon Lee
2023-10-22 00:28:29 +09:00
parent f22bc3f53c
commit b22d5521b7
4 changed files with 64 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/olive.py
+21 -13
View File
@@ -8,12 +8,9 @@ from typing import Union, Optional, Callable, List
from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection
from installer import log, args
from modules.shared import opts, cmd_opts
from modules.paths import models_path, sd_configs_path
from modules.paths import sd_configs_path
from modules.sd_models import CheckpointInfo
temp_dir = os.path.join(models_path, "OliveTemp")
cache_dir = os.path.join(models_path, "OliveCache")
submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",)
execution_provider = "CUDAExecutionProvider"
@@ -25,7 +22,18 @@ provider = (execution_provider, {
"device_id": int(cmd_opts.device_id or 0),
})
class OnnxRuntimeModel(diffusers.OnnxRuntimeModel):
config = {}
def named_modules(self):
return ()
diffusers.OnnxRuntimeModel = OnnxRuntimeModel
class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline):
model_type: str
sd_model_hash: str
sd_checkpoint_info: CheckpointInfo
sd_model_checkpoint: str
@@ -192,9 +200,9 @@ class OlivePipeline(diffusers.DiffusionPipeline):
self.original_filename = os.path.basename(path)
self.unoptimized = pipeline
del pipeline
if not os.path.exists(temp_dir):
os.mkdir(temp_dir)
self.unoptimized.save_pretrained(temp_dir)
if not os.path.exists(opts.olive_temp_dir):
os.mkdir(opts.olive_temp_dir)
self.unoptimized.save_pretrained(opts.olive_temp_dir)
@staticmethod
def from_pretrained(pretrained_model_name_or_path, **kwargs):
@@ -218,7 +226,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
if width != height:
log.warning("Olive received different width and height. The quality of the result is not guaranteed.")
out_dir = os.path.join(cache_dir, f"{self.original_filename}-{width}w-{height}h")
out_dir = os.path.join(opts.olive_cached_models_path, f"{self.original_filename}-{width}w-{height}h")
if os.path.isdir(out_dir):
del self.unoptimized
return OnnxStableDiffusionPipeline.from_pretrained(out_dir, provider=provider).apply(self)
@@ -226,7 +234,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
try:
if opts.olive_cache_optimized:
shutil.copytree(
temp_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt")
opts.olive_temp_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt")
)
optimize_config["width"] = width
@@ -260,7 +268,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
).model_path
log.info(f"Optimized {submodel}")
shutil.rmtree(temp_dir)
shutil.rmtree(opts.olive_temp_dir)
kwargs = {
"tokenizer": self.unoptimized.tokenizer,
@@ -295,9 +303,9 @@ class OlivePipeline(diffusers.DiffusionPipeline):
if os.path.isfile(weights_src_path):
weights_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data"))
shutil.copyfile(weights_src_path, weights_dst_path)
except Exception as e:
except Exception:
log.error(f"Failed to optimize model '{self.original_filename}'.")
shutil.rmtree(temp_dir, ignore_errors=True)
shutil.rmtree(opts.olive_temp_dir, ignore_errors=True)
shutil.rmtree(out_dir, ignore_errors=True)
pipeline = None
shutil.rmtree("cache", ignore_errors=True)
@@ -312,7 +320,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
optimize_config = {
"is_sdxl": False,
"source": os.path.abspath(temp_dir),
"source": os.path.abspath(opts.olive_temp_dir),
"width": 512,
"height": 512,
modules/sd_models.py
+33
View File
@@ -147,6 +147,7 @@ def list_models():
model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"]))
if shared.backend == shared.Backend.DIFFUSERS:
model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir, clear=True)
model_list += modelloader.load_diffusers_models(model_path=shared.opts.olive_sideloaded_models_path, command_path=shared.opts.olive_sideloaded_models_path)
for filename in sorted(model_list, key=str.lower):
checkpoint_info = CheckpointInfo(filename)
if checkpoint_info.name is not None:
@@ -790,6 +791,38 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"')
pipeline, model_type = detect_pipeline(checkpoint_info.path, op)
if os.path.isdir(checkpoint_info.path):
if shared.opts.olive_sideloaded_models_path in checkpoint_info.path:
try:
from modules.olive import OnnxStableDiffusionPipeline, provider
sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path, provider=provider)
sd_model.model_type = sd_model.__class__.__name__
except Exception as e:
shared.log.error(f'Failed loading {op}: {checkpoint_info.path} olive={e}')
return
else:
err1 = None
err2 = None
err3 = None
try: # try autopipeline first, best choice but not all pipelines are available
sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
sd_model.model_type = sd_model.__class__.__name__
except Exception as e:
err1 = e
try: # try diffusion pipeline next second-best choice, works for most non-linked pipelines
if err1 is not None:
sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
sd_model.model_type = sd_model.__class__.__name__
except Exception as e:
err2 = e
try: # try basic pipeline next just in case
if err2 is not None:
sd_model = diffusers.StableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
sd_model.model_type = sd_model.__class__.__name__
except Exception as e:
err3 = e # ignore last error
if err3 is not None:
shared.log.error(f'Failed loading {op}: {checkpoint_info.path} auto={err1} diffusion={err2}')
return
if model_type in ['InstaFlow']: # forced pipeline
sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
else:
modules/shared.py
+3
View File
@@ -455,6 +455,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), {
"embeddings_dir": OptionInfo(os.path.join(paths.models_path, 'embeddings'), "Folder with textual inversion embeddings", folder=True),
"hypernetwork_dir": OptionInfo(os.path.join(paths.models_path, 'hypernetworks'), "Folder with Hypernetwork models", folder=True),
"control_dir": OptionInfo(os.path.join(paths.models_path, 'control'), "Folder with Control models", folder=True),
"olive_temp_dir": OptionInfo(os.path.join(paths.models_path, 'Olive', 'temp'), "Directory for olive optimization process", folder=True),
"codeformer_models_path": OptionInfo(os.path.join(paths.models_path, 'Codeformer'), "Folder with codeformer models", folder=True),
"gfpgan_models_path": OptionInfo(os.path.join(paths.models_path, 'GFPGAN'), "Folder with GFPGAN models", folder=True),
"esrgan_models_path": OptionInfo(os.path.join(paths.models_path, 'ESRGAN'), "Folder with ESRGAN models", folder=True),
@@ -464,6 +465,8 @@ options_templates.update(options_section(('system-paths', "System Paths"), {
"swinir_models_path": OptionInfo(os.path.join(paths.models_path, 'SwinIR'), "Folder with SwinIR models", folder=True),
"ldsr_models_path": OptionInfo(os.path.join(paths.models_path, 'LDSR'), "Folder with LDSR models", folder=True),
"clip_models_path": OptionInfo(os.path.join(paths.models_path, 'CLIP'), "Folder with CLIP models", folder=True),
"olive_cached_models_path": OptionInfo(os.path.join(paths.models_path, 'Olive', 'cache'), "Folder with olive optimized cached models", folder=True),
"olive_sideloaded_models_path": OptionInfo(os.path.join(paths.models_path, 'Olive', 'sideloaded'), "Folder with olive optimized sideloaded models", folder=True),
"other_paths_sep_options": OptionInfo("<h2>Other paths</h2>", "", gr.HTML),
"openvino_cache_path": OptionInfo('cache', "Directory for OpenVINO cache", folder=True),
modules/ui_models.py
+7 -6
View File
@@ -373,10 +373,10 @@ def create_ui():
def hf_select(evt: gr.SelectData, data):
return data[evt.index[0]][0]
def hf_download_model(hub_id: str, token, variant, revision, mirror, custom_pipeline):
def hf_download_model(hub_id: str, token, variant, revision, mirror, olive_optimized):
from modules.modelloader import download_diffusers_model
download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror, custom_pipeline=custom_pipeline)
from modules.sd_models import list_models # pylint: disable=W0621
download_diffusers_model(hub_id, cache_dir=opts.olive_sideloaded_models_path if olive_optimized else opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror)
from modules.sd_models import list_models # pylint: disable=W0621
list_models()
log.info(f'Diffuser model downloaded: model="{hub_id}"')
return f'Diffuser model downloaded: model="{hub_id}"'
@@ -392,8 +392,9 @@ def create_ui():
hf_selected = gr.Textbox('', label='Select model', placeholder='select model from search results or enter model name manually')
with gr.Column(scale=1):
with gr.Row():
hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label='Specify model variant', placeholder='')
hf_revision = gr.Textbox('', label='Specify model revision', placeholder='')
hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label = 'Specify model variant', placeholder='')
hf_revision = gr.Textbox('', label = 'Specify model revision', placeholder='')
hf_olive = gr.Checkbox(False, label = 'Olive optimized')
with gr.Row():
hf_token = gr.Textbox('', label='Huggingface token', placeholder='optional access token for private or gated models')
hf_mirror = gr.Textbox('', label='Huggingface mirror', placeholder='optional mirror site for downloads')
@@ -410,7 +411,7 @@ def create_ui():
hf_search_text.submit(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results])
hf_search_btn.click(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results])
hf_results.select(fn=hf_select, inputs=[hf_results], outputs=[hf_selected])
hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_custom_pipeline], outputs=[models_outcome])
hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_olive], outputs=[models_outcome])
with gr.Tab(label="CivitAI"):
data = []