redesign live preview and bring full quality toggle to original backend

This commit is contained in:
Vladimir Mandic
2023-10-16 14:07:00 -04:00
parent 21595ee9e7
commit bb912c35f9
8 changed files with 106 additions and 63 deletions
+23 -14
View File
@@ -35,6 +35,7 @@ import modules.sd_samplers_common
import modules.sd_models
import modules.sd_vae
import modules.sd_vae_approx
import modules.taesd.sd_vae_taesd
import modules.generation_parameters_copypaste
from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet, hypertile_set
@@ -433,15 +434,25 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
return x
def decode_first_stage(model, x):
def decode_first_stage(model, x, full_quality=True):
with devices.autocast(disable = x.dtype==devices.dtype_vae):
if hasattr(model, 'decode_first_stage'):
x = model.decode_first_stage(x)
elif hasattr(model, 'vae'):
x = model.vae(x)
else:
shared.log.warning('Cannot decode first stage')
return x
try:
if full_quality:
if hasattr(model, 'decode_first_stage'):
x_sample = model.decode_first_stage(x)
elif hasattr(model, 'vae'):
x_sample = model.vae(x)
else:
x_sample = x
shared.log.error('Decode VAE unknown model')
else:
x_sample = torch.zeros((len(x), 3, x.shape[2] * 8, x.shape[3] * 8), dtype=devices.dtype_vae, device=devices.device)
for i in range(len(x_sample)):
x_sample[i] = (modules.taesd.sd_vae_taesd.decode(x[i]) * 2.0) - 1.0
except Exception as e:
x_sample = x
shared.log.error(f'Decode VAE: {e}')
return x_sample
def get_fixed_seed(seed):
@@ -755,8 +766,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
t0 = time.time()
with devices.autocast():
p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
if shared.opts.live_previews_enable and shared.opts.show_progress_type == "Approximate NN" and shared.backend == shared.Backend.ORIGINAL:
modules.sd_vae_approx.model()
if shared.state.job_count == -1:
shared.state.job_count = p.n_iter
extra_network_data = None
@@ -801,7 +810,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
comments[comment] = 1
with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():
samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts)
x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae), p.full_quality)[0].cpu() for i in range(samples_ddim.size(0))]
try:
for x in x_samples_ddim:
devices.test_for_nans(x, "vae")
@@ -811,7 +820,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
devices.dtype_vae = torch.bfloat16
vae_file, vae_source = modules.sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint)
modules.sd_vae.load_vae(p.sd_model, vae_file, vae_source)
x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae), p.full_quality)[0].cpu() for i in range(samples_ddim.size(0))]
for x in x_samples_ddim:
devices.test_for_nans(x, "vae")
else:
@@ -1067,7 +1076,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
for i in range(samples.shape[0]):
save_intermediate(samples, i)
if latent_scale_mode is None or self.hr_force: # non-latent upscaling
decoded_samples = decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae))
decoded_samples = decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae), self.full_quality)
lowres_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)
batch_images = []
for _i, x_sample in enumerate(lowres_samples):
@@ -1093,7 +1102,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
else:
samples = torch.nn.functional.interpolate(samples, size=(target_height // 8, target_width // 8), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"])
if getattr(self, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0:
image_conditioning = self.img2img_image_conditioning(decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae)), samples)
image_conditioning = self.img2img_image_conditioning(decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae), self.full_quality), samples)
else:
image_conditioning = self.txt2img_image_conditioning(samples.to(dtype=devices.dtype_vae))
if self.latent_sampler == "PLMS":