PromptForge 1.29.8

preview_decoder.py

@@ -0,0 +1,78 @@
+"""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