Austin Mroz
a0c5c53f37
Load Images, Load Video, Video Combine, and Batched VAE nodes have had progress bars added. For Load Video, a progress bar is only shown if frame_load_cap is set and if the video has fewer frames than that cap, the length of the progress bar may not be accurate.
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import torch
|
|
from nodes import VAEEncode
|
|
from comfy.utils import ProgressBar
|
|
|
|
|
|
class VAEDecodeBatched:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT", ),
|
|
"vae": ("VAE", ),
|
|
"per_batch": ("INT", {"default": 16, "min": 1})
|
|
}
|
|
}
|
|
|
|
CATEGORY = "Video Helper Suite 🎥🅥🅗🅢/batched nodes"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "decode"
|
|
|
|
def decode(self, vae, samples, per_batch):
|
|
decoded = []
|
|
pbar = ProgressBar(samples["samples"].shape[0])
|
|
for start_idx in range(0, samples["samples"].shape[0], per_batch):
|
|
decoded.append(vae.decode(samples["samples"][start_idx:start_idx+per_batch]))
|
|
pbar.update(per_batch)
|
|
return (torch.cat(decoded, dim=0), )
|
|
|
|
|
|
class VAEEncodeBatched:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pixels": ("IMAGE", ), "vae": ("VAE", ),
|
|
"per_batch": ("INT", {"default": 16, "min": 1})
|
|
}
|
|
}
|
|
|
|
CATEGORY = "Video Helper Suite 🎥🅥🅗🅢/batched nodes"
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "encode"
|
|
|
|
def encode(self, vae, pixels, per_batch):
|
|
t = []
|
|
pbar = ProgressBar(pixels.shape[0])
|
|
for start_idx in range(0, pixels.shape[0], per_batch):
|
|
try:
|
|
sub_pixels = vae.vae_encode_crop_pixels(pixels[start_idx:start_idx+per_batch])
|
|
except:
|
|
sub_pixels = VAEEncode.vae_encode_crop_pixels(pixels[start_idx:start_idx+per_batch])
|
|
t.append(vae.encode(sub_pixels[:,:,:,:3]))
|
|
pbar.update(per_batch)
|
|
return ({"samples": torch.cat(t, dim=0)}, )
|