Added some features for an LCM condenser plugin

toolkit/image_utils.py

@@ -5,9 +5,12 @@ import json
 import os
 import io
 import struct
+from typing import TYPE_CHECKING
 
 import cv2
 import numpy as np
+import torch
+from diffusers import AutoencoderTiny
 
 FILE_UNKNOWN = "Sorry, don't know how to get size for this file."
 
@@ -424,23 +427,47 @@ def main(argv=None):
 is_window_shown = False
 
 
-def show_img(img):
+def show_img(img, name='AI Toolkit'):
     global is_window_shown
 
     img = np.clip(img, 0, 255).astype(np.uint8)
-    cv2.imshow('AI Toolkit', img[:, :, ::-1])
+    cv2.imshow(name, img[:, :, ::-1])
     k = cv2.waitKey(10) & 0xFF
     if k == 27:  # Esc key to stop
         print('\nESC pressed, stopping')
         raise KeyboardInterrupt
-    # show again to initialize the window if first
     if not is_window_shown:
-        cv2.imshow('AI Toolkit', img[:, :, ::-1])
-        k = cv2.waitKey(10) & 0xFF
-        if k == 27:  # Esc key to stop
-            print('\nESC pressed, stopping')
-            raise KeyboardInterrupt
-    is_window_shown = True
+        is_window_shown = True
+
+
+
+def show_tensors(imgs: torch.Tensor, name='AI Toolkit'):
+    # if rank is 4
+    if len(imgs.shape) == 4:
+        img_list = torch.chunk(imgs, imgs.shape[0], dim=0)
+    else:
+        img_list = [imgs]
+    # put images side by side
+    img = torch.cat(img_list, dim=3)
+    # img is -1 to 1, convert to 0 to 255
+    img = img / 2 + 0.5
+    img_numpy = img.to(torch.float32).detach().cpu().numpy()
+    img_numpy = np.clip(img_numpy, 0, 1) * 255
+    # convert to numpy Move channel to last
+    img_numpy = img_numpy.transpose(0, 2, 3, 1)
+    # convert to uint8
+    img_numpy = img_numpy.astype(np.uint8)
+    show_img(img_numpy[0], name=name)
+
+
+def show_latents(latents: torch.Tensor, vae: 'AutoencoderTiny', name='AI Toolkit'):
+    # decode latents
+    if vae.device == 'cpu':
+        vae.to(latents.device)
+    latents = latents / vae.config['scaling_factor']
+    imgs = vae.decode(latents).sample
+    show_tensors(imgs, name=name)
+
 
 
 def on_exit():