Update inference.py

inference.py

@@ -19,6 +19,7 @@ import math
 import os
 import sys
 import time
+from tqdm import tqdm
 
 import torch
 import torch.nn as nn
@@ -359,7 +360,7 @@ def generate(args: argparse.Namespace, device: torch.device) -> None:
     log(f"[sample] Euler flow, {args.steps} steps ...")
     t_sample = time.perf_counter()
 
-    for i in range(args.steps):
+    for i in tqdm(range(args.steps), desc="Generating", unit="step", dynamic_ncols=True):
         t = torch.full((args.n,), i * dt, device=device)
         with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_amp):
             if use_cfg:
@@ -369,8 +370,6 @@ def generate(args: argparse.Namespace, device: torch.device) -> None:
             else:
                 v = model(z, t, ctx, cmask).float()
         z = z + dt * v
-        if (i + 1) % max(1, args.steps // 10) == 0 or i == 0:
-            log(f"  step {i + 1}/{args.steps}")
 
     log(f"[sample] denoising done in {time.perf_counter() - t_sample:.1f}s")
     log("[vae] decoding latents ...")