update image saving tool

src/utils/util.py

@@ -82,6 +82,19 @@ def save_videos_from_pil(pil_images, path, fps=8):
     else:
         raise ValueError("Unsupported file type. Use .mp4 or .gif.")
 
+def save_pil_imgs(videos: torch.Tensor, path: str, rescale=False):
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    os.makedirs(path, exist_ok=True)
+
+    for idx, x in enumerate(videos):
+        x = torchvision.utils.make_grid(x, nrow=1)  # (c h w)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)  # (h w c)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = (x * 255).numpy().astype(np.uint8)
+        x = Image.fromarray(x)
+        x.save(os.path.join(path, f"{idx:05d}.png"))
+        
 
 def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
     videos = rearrange(videos, "b c t h w -> t b c h w")