fix taesd live preview

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-10-27 09:23:51 -04:00
parent c9021ea7d3
commit 706852e7c5
2 changed files with 8 additions and 4 deletions
+6 -2
View File
@@ -1,5 +1,6 @@
import io
import sys
import time
import json
import copy
import inspect
@@ -453,7 +454,6 @@ def move_model(model, device=None, force=False):
if hasattr(model.vae, '_hf_hook'):
debug_move(f'Model move: to={device} class={model.vae.__class__} fn={fn}') # pylint: disable=protected-access
model.vae._hf_hook.execution_device = device # pylint: disable=protected-access
debug_move(f'Model move: device={device} class={model.__class__} accelerate={getattr(model, "has_accelerate", False)} fn={fn}') # pylint: disable=protected-access
if hasattr(model, "components"): # accelerate patch
for name, m in model.components.items():
if not hasattr(m, "_hf_hook"): # not accelerate hook
@@ -472,8 +472,9 @@ def move_model(model, device=None, force=False):
if hasattr(model, "device") and devices.normalize_device(model.device) == devices.normalize_device(device):
return
try:
t0 = time.time()
try:
model.to(device)
model.to(device, non_blocking=True)
if hasattr(model, "prior_pipe"):
model.prior_pipe.to(device)
except Exception as e0:
@@ -493,8 +494,11 @@ def move_model(model, device=None, force=False):
pass # ignore model move if sequential offload is enabled
else:
raise e0
t1 = time.time()
except Exception as e1:
shared.log.error(f'Model move: device={device} {e1}')
if os.environ.get('SD_MOVE_DEBUG', None) or (t1-t0) > 0.1:
shared.log.debug(f'Model move: device={device} class={model.__class__.__name__} accelerate={getattr(model, "has_accelerate", False)} fn={fn} time={t1-t0:.2f}') # pylint: disable=protected-access
devices.torch_gc()
+2 -2
View File
@@ -171,12 +171,12 @@ def decode(latents):
try:
with devices.inference_context():
latents = latents.detach().clone().to(devices.device, dtype)
if len(latents.shape) == 3 and latents.shape[0] == 4:
if len(latents.shape) == 3:
latents = latents.unsqueeze(0)
image = vae.decoder(latents).clamp(0, 1).detach()
image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
return image[0]
elif len(latents.shape) == 4 and latents.shape[1] == 4:
elif len(latents.shape) == 4:
image = vae.decoder(latents).clamp(0, 1).detach()
image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
return image