79 lines
3.3 KiB
Python
79 lines
3.3 KiB
Python
"""Tiny decoders for live previews: the small TAESD-style decoders in models/vae_approx (taeh3 for MiniMax H3).
|
|
|
|
These files are flat decoders whose keys are the positions of their modules (`N.conv.0.weight` is a block, `N.weight`
|
|
a convolution, a missing number is a module without weights), so the network is rebuilt from the file itself with
|
|
ComfyUI's own TAESD blocks. Each latent frame is decoded on its own, which is what a preview needs.
|
|
"""
|
|
import logging
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import comfy.model_management as model_management
|
|
import comfy.utils
|
|
import folder_paths
|
|
from comfy.taesd.taesd import Block, Clamp, conv
|
|
|
|
LATENT2RGB = "latent2rgb (instant, no file)"
|
|
|
|
|
|
def decoder_names():
|
|
return [LATENT2RGB] + folder_paths.get_filename_list("vae_approx")
|
|
|
|
|
|
def default_decoder():
|
|
return next((name for name in folder_paths.get_filename_list("vae_approx") if name.lower().startswith("taeh3")), LATENT2RGB)
|
|
|
|
|
|
def _build(state):
|
|
groups = {}
|
|
for key, value in state.items():
|
|
head, _, rest = key.partition(".")
|
|
if not head.isdigit():
|
|
raise ValueError(f"not a flat tiny decoder (key {key})")
|
|
groups.setdefault(int(head), {})[rest] = value
|
|
modules = []
|
|
for index in range(max(groups) + 1):
|
|
entry = groups.get(index)
|
|
if entry is None:
|
|
modules.append(Clamp() if index == 0 else nn.ReLU() if index == 2 else nn.Upsample(scale_factor=2))
|
|
elif "conv.0.weight" in entry:
|
|
weight = entry["conv.0.weight"]
|
|
modules.append(Block(weight.shape[1], weight.shape[0], use_midblock_gn="pool.0.weight" in entry))
|
|
elif "weight" in entry:
|
|
weight = entry["weight"]
|
|
modules.append(conv(weight.shape[1], weight.shape[0], bias="bias" in entry))
|
|
else:
|
|
raise ValueError(f"unrecognized module {index} ({sorted(entry)})")
|
|
return nn.Sequential(*modules)
|
|
|
|
|
|
class TinyDecoder:
|
|
def __init__(self, state, device, dtype):
|
|
self.model = _build(state)
|
|
self.model.load_state_dict(state)
|
|
self.model = self.model.eval().to(device=device, dtype=dtype)
|
|
self.device, self.dtype = device, dtype
|
|
self.channels = self.model[1].weight.shape[1]
|
|
|
|
def frames(self, video, indices):
|
|
"""The latent frames `indices` of a video latent [1, C, T, h, w] as RGB in 0..1, [n, h', w', 3] on the CPU."""
|
|
out = []
|
|
for index in indices:
|
|
frame = self.model(video[0, :, index].unsqueeze(0).to(device=self.device, dtype=self.dtype))[0]
|
|
out.append(frame.clamp(0, 1).float().cpu().movedim(0, -1))
|
|
return torch.stack(out)
|
|
|
|
|
|
def load(name):
|
|
"""The tiny decoder `name` from models/vae_approx, or None (with a log line) when it is not a flat decoder or does not fit."""
|
|
path = folder_paths.get_full_path("vae_approx", name) if name != LATENT2RGB else None
|
|
if path is None:
|
|
return None
|
|
try:
|
|
device = model_management.get_torch_device()
|
|
return TinyDecoder(comfy.utils.load_torch_file(path, safe_load=True), device, model_management.vae_dtype(device, [torch.float16, torch.bfloat16]))
|
|
except (ValueError, KeyError, RuntimeError) as error:
|
|
logging.warning(f"[PromptForge] preview decoder {name} cannot be used ({error}); using latent2rgb.")
|
|
return None
|