return live preview defaults to how they were
only download TAESD model when it's needed return calculations in single_sample_to_image to just if/elif/elif blocks keep taesd model in its own directory
This commit is contained in:
parent
b217ebc490
commit
56a2672831
@ -22,28 +22,29 @@ def setup_img2img_steps(p, steps=None):
|
|||||||
return steps, t_enc
|
return steps, t_enc
|
||||||
|
|
||||||
|
|
||||||
approximation_indexes = {"Full": 0, "Tiny AE": 1, "Approx NN": 2, "Approx cheap": 3}
|
approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2, "TAESD": 3}
|
||||||
|
|
||||||
|
|
||||||
def single_sample_to_image(sample, approximation=None):
|
def single_sample_to_image(sample, approximation=None):
|
||||||
if approximation is None or approximation not in approximation_indexes.keys():
|
|
||||||
approximation = approximation_indexes.get(opts.show_progress_type, 1)
|
|
||||||
|
|
||||||
if approximation == 1:
|
if approximation is None:
|
||||||
x_sample = sd_vae_taesd.decode()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
|
approximation = approximation_indexes.get(opts.show_progress_type, 0)
|
||||||
x_sample = sd_vae_taesd.TAESD.unscale_latents(x_sample)
|
|
||||||
x_sample = torch.clamp((x_sample * 0.25) + 0.5, 0, 1)
|
if approximation == 2:
|
||||||
|
x_sample = sd_vae_approx.cheap_approximation(sample)
|
||||||
|
elif approximation == 1:
|
||||||
|
x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
|
||||||
|
elif approximation == 3:
|
||||||
|
x_sample = sd_vae_taesd.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
|
||||||
|
x_sample = sd_vae_taesd.TAESD.unscale_latents(x_sample) # returns value in [-2, 2]
|
||||||
|
x_sample = x_sample * 0.5
|
||||||
else:
|
else:
|
||||||
if approximation == 3:
|
x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
|
||||||
x_sample = sd_vae_approx.cheap_approximation(sample)
|
|
||||||
elif approximation == 2:
|
|
||||||
x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
|
|
||||||
else:
|
|
||||||
x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
|
|
||||||
x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
|
|
||||||
|
|
||||||
|
x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
|
||||||
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
|
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
|
||||||
x_sample = x_sample.astype(np.uint8)
|
x_sample = x_sample.astype(np.uint8)
|
||||||
|
|
||||||
return Image.fromarray(x_sample)
|
return Image.fromarray(x_sample)
|
||||||
|
|
||||||
|
|
||||||
|
@ -61,16 +61,28 @@ class TAESD(nn.Module):
|
|||||||
return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
|
return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
|
||||||
|
|
||||||
|
|
||||||
def decode():
|
def download_model(model_path):
|
||||||
|
model_url = 'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth'
|
||||||
|
|
||||||
|
if not os.path.exists(model_path):
|
||||||
|
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
||||||
|
|
||||||
|
print(f'Downloading TAESD decoder to: {model_path}')
|
||||||
|
torch.hub.download_url_to_file(model_url, model_path)
|
||||||
|
|
||||||
|
|
||||||
|
def model():
|
||||||
global sd_vae_taesd
|
global sd_vae_taesd
|
||||||
|
|
||||||
if sd_vae_taesd is None:
|
if sd_vae_taesd is None:
|
||||||
model_path = os.path.join(paths_internal.models_path, "VAE-approx", "taesd_decoder.pth")
|
model_path = os.path.join(paths_internal.models_path, "VAE-taesd", "taesd_decoder.pth")
|
||||||
|
download_model(model_path)
|
||||||
|
|
||||||
if os.path.exists(model_path):
|
if os.path.exists(model_path):
|
||||||
sd_vae_taesd = TAESD(model_path)
|
sd_vae_taesd = TAESD(model_path)
|
||||||
sd_vae_taesd.eval()
|
sd_vae_taesd.eval()
|
||||||
sd_vae_taesd.to(devices.device, devices.dtype)
|
sd_vae_taesd.to(devices.device, devices.dtype)
|
||||||
else:
|
else:
|
||||||
raise FileNotFoundError('Tiny AE model not found')
|
raise FileNotFoundError('TAESD model not found')
|
||||||
|
|
||||||
return sd_vae_taesd.decoder
|
return sd_vae_taesd.decoder
|
||||||
|
@ -425,7 +425,7 @@ options_templates.update(options_section(('ui', "Live previews"), {
|
|||||||
"live_previews_enable": OptionInfo(True, "Show live previews of the created image"),
|
"live_previews_enable": OptionInfo(True, "Show live previews of the created image"),
|
||||||
"show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),
|
"show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),
|
||||||
"show_progress_every_n_steps": OptionInfo(10, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
|
"show_progress_every_n_steps": OptionInfo(10, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
|
||||||
"show_progress_type": OptionInfo("Tiny AE", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Tiny AE", "Approx NN", "Approx cheap"]}),
|
"show_progress_type": OptionInfo("Approx NN", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap", "TAESD"]}),
|
||||||
"live_preview_content": OptionInfo("Prompt", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"]}),
|
"live_preview_content": OptionInfo("Prompt", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"]}),
|
||||||
"live_preview_refresh_period": OptionInfo(1000, "Progressbar/preview update period, in milliseconds")
|
"live_preview_refresh_period": OptionInfo(1000, "Progressbar/preview update period, in milliseconds")
|
||||||
}))
|
}))
|
||||||
|
11
webui.py
11
webui.py
@ -144,21 +144,10 @@ Use --skip-version-check commandline argument to disable this check.
|
|||||||
""".strip())
|
""".strip())
|
||||||
|
|
||||||
|
|
||||||
def check_taesd():
|
|
||||||
from modules.paths_internal import models_path
|
|
||||||
|
|
||||||
model_url = 'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth'
|
|
||||||
model_path = os.path.join(models_path, "VAE-approx", "taesd_decoder.pth")
|
|
||||||
if not os.path.exists(model_path):
|
|
||||||
print('From taesd repo download decoder model')
|
|
||||||
torch.hub.download_url_to_file(model_url, model_path)
|
|
||||||
|
|
||||||
|
|
||||||
def initialize():
|
def initialize():
|
||||||
fix_asyncio_event_loop_policy()
|
fix_asyncio_event_loop_policy()
|
||||||
|
|
||||||
check_versions()
|
check_versions()
|
||||||
check_taesd()
|
|
||||||
|
|
||||||
extensions.list_extensions()
|
extensions.list_extensions()
|
||||||
localization.list_localizations(cmd_opts.localizations_dir)
|
localization.list_localizations(cmd_opts.localizations_dir)
|
||||||
|
Loading…
Reference in New Issue
Block a user