This commit is contained in:
lvmin 2023-08-10 08:54:51 -07:00
parent 265c30a4ce
commit 3bafd2ea6c
2 changed files with 4 additions and 3 deletions

View File

@ -58,13 +58,14 @@ def get_previewer(device, latent_format):
taesd = TAESD(None, taesd_decoder_path).to(device)
def preview_function(x0):
def preview_function(x0, step, total_steps):
with torch.no_grad():
x_sample = taesd.decoder(x0).detach() * 255.0
x_sample = einops.rearrange(x_sample, 'b c h w -> b h w c')
x_sample = x_sample.cpu().numpy()[..., ::-1].copy().clip(0, 255).astype(np.uint8)
for i, s in enumerate(x_sample):
cv2.imshow(f'Preview {i}', s)
cv2.setWindowTitle(f'Preview {i}', f'Preview {i}, [{step}/{total_steps}]')
cv2.waitKey(1)
taesd.preview = preview_function
@ -99,7 +100,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=9.0, sa
def callback(step, x0, x, total_steps):
if previewer and step % 3 == 0:
previewer.preview(x0)
previewer.preview(x0, step, total_steps)
pbar.update_absolute(step + 1, total_steps, None)
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,

View File

@ -44,5 +44,5 @@ def process(positive_prompt, negative_prompt, width=1024, height=1024, batch_siz
images = core.image_to_numpy(decoded_latent)
core.close_all_preview()
return images