From e49105ece5849e43a69241da136e62fc9b5d1912 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Tue, 20 Aug 2024 14:52:05 +0000 Subject: [PATCH] Add tqdm progress bar to FLUX denoising. --- invokeai/backend/flux/sampling.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/invokeai/backend/flux/sampling.py b/invokeai/backend/flux/sampling.py index 3e3c933d4e..5d670c3e69 100644 --- a/invokeai/backend/flux/sampling.py +++ b/invokeai/backend/flux/sampling.py @@ -4,6 +4,7 @@ from typing import Callable import torch from einops import rearrange, repeat from torch import Tensor +from tqdm import tqdm from .model import Flux from .modules.conditioner import HFEncoder @@ -115,7 +116,7 @@ def denoise( # this is ignored for schnell guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) - for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:], strict=True): + for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:], strict=True))): t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) pred = model( img=img,