104 lines
2.6 KiB
Python
104 lines
2.6 KiB
Python
from collections import defaultdict
|
|
from typing import Dict, List
|
|
import torch
|
|
from tqdm import trange
|
|
from .model.iter import try_get_iter
|
|
|
|
class VAEDecodeBatched:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT", ),
|
|
"vae": ("VAE", ),
|
|
"batch_size": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 32,
|
|
"step": 1
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "decode"
|
|
|
|
CATEGORY = "latent"
|
|
|
|
def decode(self, vae, samples, batch_size: int):
|
|
s = samples['samples']
|
|
n = s.shape[0]
|
|
|
|
iters = try_get_iter(vae)
|
|
if iters is None:
|
|
vae_num = 1
|
|
else:
|
|
vae_num = len(iters)
|
|
|
|
vae_results: Dict[int,List[torch.Tensor]] = defaultdict(lambda: [])
|
|
|
|
for i in trange(0, n, batch_size):
|
|
e = min([i+batch_size, n])
|
|
t = s[i:e, ...]
|
|
v = vae.decode(t)
|
|
|
|
vaes = torch.chunk(v, vae_num)
|
|
|
|
for vn, vv in enumerate(vaes):
|
|
vae_results[vn].append(vv)
|
|
|
|
results = []
|
|
for k in sorted(vae_results.keys()):
|
|
v = vae_results[k]
|
|
results.extend(v)
|
|
|
|
vs = torch.cat(results).contiguous()
|
|
return (vs,)
|
|
|
|
|
|
class VAEEncodeBatched:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pixels": ("IMAGE", ),
|
|
"vae": ("VAE", ),
|
|
"batch_size": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 32,
|
|
"step": 1
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "latent"
|
|
|
|
def encode(self, vae, pixels, batch_size: int):
|
|
n = pixels.shape[0]
|
|
x = (pixels.shape[1] // 64) * 64
|
|
y = (pixels.shape[2] // 64) * 64
|
|
if pixels.shape[1] != x or pixels.shape[2] != y:
|
|
pixels = pixels[:,:x,:y,:]
|
|
|
|
pixels = pixels[:,:,:,:3]
|
|
|
|
results = []
|
|
for i in trange(0, n, batch_size):
|
|
e = max([i+batch_size, n])
|
|
t = pixels[i:e, ...]
|
|
v = vae.encode(t)
|
|
results.append(v)
|
|
|
|
vs = torch.cat(results)
|
|
return ({"samples":vs}, )
|