Fix Full VAE previews

This commit is contained in:
Disty0
2024-07-28 02:03:53 +03:00
parent 4492dedf16
commit fa1e77cc6d
2 changed files with 5 additions and 5 deletions
+3 -3
View File
@@ -196,9 +196,9 @@ def decode_first_stage(model, x, full_quality=True):
try:
if full_quality:
if hasattr(model, 'decode_first_stage'):
x_sample = model.decode_first_stage(x)
x_sample = model.decode_first_stage(x) * 0.5 + 0.5
elif hasattr(model, 'vae'):
x_sample = model.vae(x)
x_sample = processing_vae.vae_decode(latents=x, model=model, output_type='np', full_quality=full_quality)
else:
x_sample = x
shared.log.error('Decode VAE unknown model')
@@ -206,7 +206,7 @@ def decode_first_stage(model, x, full_quality=True):
from modules import sd_vae_taesd
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] = sd_vae_taesd.decode(x[i])
x_sample[i] = sd_vae_taesd.decode(x[i]) * 0.5 + 0.5
except Exception as e:
x_sample = x
shared.log.error(f'Decode VAE: {e}')
+2 -2
View File
@@ -61,7 +61,7 @@ def single_sample_to_image(sample, approximation=None):
if approximation == 2: # TAESD
x_sample = sd_vae_taesd.decode(sample)
x_sample = (1.0 + x_sample) / 2.0 # preview requires smaller range
elif sd_cascade:
elif sd_cascade and not approximation == 3:
x_sample = sd_vae_stablecascade.decode(sample)
elif approximation == 0: # Simple
x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
@@ -70,7 +70,7 @@ def single_sample_to_image(sample, approximation=None):
if shared.sd_model_type == "sdxl":
x_sample = x_sample[[2,1,0], :, :] # BGR to RGB
elif approximation == 3: # Full VAE
x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5
x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
else:
warn_once(f"Unknown latent decode type: {approximation}")
return Image.new(mode="RGB", size=(512, 512))