mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2024-06-07 21:20:49 +00:00
LDSR cache / optimization / opt_channelslast
This commit is contained in:
parent
7cea280a8f
commit
6df316c881
@ -11,25 +11,41 @@ from omegaconf import OmegaConf
|
|||||||
|
|
||||||
from ldm.models.diffusion.ddim import DDIMSampler
|
from ldm.models.diffusion.ddim import DDIMSampler
|
||||||
from ldm.util import instantiate_from_config, ismap
|
from ldm.util import instantiate_from_config, ismap
|
||||||
|
from modules import shared, sd_hijack
|
||||||
|
|
||||||
warnings.filterwarnings("ignore", category=UserWarning)
|
warnings.filterwarnings("ignore", category=UserWarning)
|
||||||
|
|
||||||
|
cached_ldsr_model: torch.nn.Module = None
|
||||||
|
|
||||||
|
|
||||||
# Create LDSR Class
|
# Create LDSR Class
|
||||||
class LDSR:
|
class LDSR:
|
||||||
def load_model_from_config(self, half_attention):
|
def load_model_from_config(self, half_attention):
|
||||||
print(f"Loading model from {self.modelPath}")
|
global cached_ldsr_model
|
||||||
pl_sd = torch.load(self.modelPath, map_location="cpu")
|
|
||||||
sd = pl_sd["state_dict"]
|
if shared.opts.ldsr_cached and cached_ldsr_model is not None:
|
||||||
config = OmegaConf.load(self.yamlPath)
|
print(f"Loading model from cache")
|
||||||
config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1"
|
model: torch.nn.Module = cached_ldsr_model
|
||||||
model = instantiate_from_config(config.model)
|
else:
|
||||||
model.load_state_dict(sd, strict=False)
|
print(f"Loading model from {self.modelPath}")
|
||||||
model.cuda()
|
pl_sd = torch.load(self.modelPath, map_location="cpu")
|
||||||
if half_attention:
|
sd = pl_sd["state_dict"]
|
||||||
model = model.half()
|
config = OmegaConf.load(self.yamlPath)
|
||||||
|
config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1"
|
||||||
|
model: torch.nn.Module = instantiate_from_config(config.model)
|
||||||
|
model.load_state_dict(sd, strict=False)
|
||||||
|
model = model.to(shared.device)
|
||||||
|
if half_attention:
|
||||||
|
model = model.half()
|
||||||
|
if shared.cmd_opts.opt_channelslast:
|
||||||
|
model = model.to(memory_format=torch.channels_last)
|
||||||
|
|
||||||
|
sd_hijack.model_hijack.hijack(model) # apply optimization
|
||||||
|
model.eval()
|
||||||
|
|
||||||
|
if shared.opts.ldsr_cached:
|
||||||
|
cached_ldsr_model = model
|
||||||
|
|
||||||
model.eval()
|
|
||||||
return {"model": model}
|
return {"model": model}
|
||||||
|
|
||||||
def __init__(self, model_path, yaml_path):
|
def __init__(self, model_path, yaml_path):
|
||||||
|
@ -59,6 +59,7 @@ def on_ui_settings():
|
|||||||
import gradio as gr
|
import gradio as gr
|
||||||
|
|
||||||
shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('upscaling', "Upscaling")))
|
shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('upscaling', "Upscaling")))
|
||||||
|
shared.opts.add_option("ldsr_cached", shared.OptionInfo(False, "Cache LDSR model in memory", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling")))
|
||||||
|
|
||||||
|
|
||||||
script_callbacks.on_ui_settings(on_ui_settings)
|
script_callbacks.on_ui_settings(on_ui_settings)
|
||||||
|
Loading…
Reference in New Issue
Block a user