- GNM Head Viewport: all-in-one node with live three.js viewport (slider changes re-pose the mesh via /gnm/viewport/* endpoints, mouse orbit/zoom syncs back to camera widgets) - GNM Head Render: pose + render to IMAGE / normal / depth / mask - Identity/Expression samplers: TF-free numpy port of GNM's Keras decoders (verified identical to tf.keras output) - pyrender-based renderer with clean-subprocess fallback for GL state pollution inside ComfyUI - Vendored gnm.shape core (numpy backend) + model data, Apache-2.0 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
"""TensorFlow-free GNM semantic samplers.
|
|
|
|
GNM's identity/expression decoders are small Keras MLPs
|
|
(concat(latent, one-hot) -> Dense relu x4 -> Dense linear). We read the
|
|
weights straight from the bundled .h5 files with h5py and run the forward
|
|
pass in numpy, so the sampler nodes do not need TensorFlow.
|
|
|
|
Verified to produce outputs identical to tf.keras (see repo tests).
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
_DATA_DIR = os.path.join(
|
|
os.path.dirname(os.path.abspath(__file__)),
|
|
"vendor", "gnm", "shape", "data", "semantic_sampler",
|
|
)
|
|
IDENTITY_DECODER_PATH = os.path.join(_DATA_DIR, "identity_decoder_model.h5")
|
|
EXPRESSION_DECODER_PATH = os.path.join(_DATA_DIR, "expression_decoder_model.h5")
|
|
|
|
_H5PY_HELP = (
|
|
"h5py is required for the GNM sampler nodes.\n"
|
|
"Install it into your ComfyUI python environment:\n"
|
|
" pip install h5py"
|
|
)
|
|
|
|
NUM_GENDER_CLASSES = 2
|
|
NUM_ETHNICITY_CLASSES = 4
|
|
NUM_EXPRESSION_CLASSES = 20
|
|
|
|
|
|
class _MLPDecoder:
|
|
"""Numpy forward pass for the Keras decoder MLPs bundled with GNM."""
|
|
|
|
def __init__(self, h5_path):
|
|
try:
|
|
import h5py # noqa: PLC0415
|
|
except ImportError as exc:
|
|
raise RuntimeError(_H5PY_HELP) from exc
|
|
|
|
with h5py.File(h5_path, "r") as f:
|
|
config = json.loads(f.attrs["model_config"])
|
|
layers = config["config"]["layers"]
|
|
|
|
input_shapes = {}
|
|
dense_names = []
|
|
for layer in layers:
|
|
cls = layer["class_name"]
|
|
conf = layer["config"]
|
|
if cls == "InputLayer":
|
|
shape = conf.get("batch_input_shape") or conf.get("batch_shape")
|
|
input_shapes[conf["name"]] = int(shape[-1])
|
|
elif cls == "Dense":
|
|
dense_names.append((conf["name"], conf["activation"]))
|
|
elif cls not in ("Concatenate",):
|
|
raise ValueError(f"Unsupported layer type in {h5_path}: {cls}")
|
|
|
|
self.latent_dim = input_shapes["latent_input"]
|
|
self.label_dim = input_shapes["decoder_label_input"]
|
|
|
|
weights = f["model_weights"]
|
|
self.layers = []
|
|
for name, activation in dense_names:
|
|
group = weights[name][name]
|
|
kernel = np.array(group["kernel:0"], dtype=np.float32)
|
|
bias = np.array(group["bias:0"], dtype=np.float32)
|
|
self.layers.append((kernel, bias, activation))
|
|
|
|
expected = self.latent_dim + self.label_dim
|
|
if self.layers[0][0].shape[0] != expected:
|
|
raise ValueError(
|
|
f"Decoder input mismatch in {h5_path}: first kernel expects "
|
|
f"{self.layers[0][0].shape[0]}, got {expected}"
|
|
)
|
|
|
|
def predict(self, z, labels):
|
|
"""z: (N, latent_dim), labels: (N, label_dim) -> (N, out_dim)."""
|
|
x = np.concatenate(
|
|
[np.asarray(z, np.float32), np.asarray(labels, np.float32)], axis=-1
|
|
)
|
|
for kernel, bias, activation in self.layers:
|
|
x = x @ kernel + bias
|
|
if activation == "relu":
|
|
x = np.maximum(x, 0.0)
|
|
elif activation != "linear":
|
|
raise ValueError(f"Unsupported activation: {activation}")
|
|
return x
|
|
|
|
|
|
def _one_hot(index, num_classes):
|
|
v = np.zeros((1, num_classes), dtype=np.float32)
|
|
v[0, int(index)] = 1.0
|
|
return v
|
|
|
|
|
|
class IdentitySampler:
|
|
"""Numpy port of gnm.shape.semantic_sampler.IdentitySampler."""
|
|
|
|
def __init__(self, decoder_path=IDENTITY_DECODER_PATH):
|
|
self._decoder = _MLPDecoder(decoder_path)
|
|
|
|
def sample(self, gender_index, ethnicity_index, num_samples=1, rng=None):
|
|
rng = rng or np.random.default_rng()
|
|
labels = np.concatenate(
|
|
[
|
|
_one_hot(gender_index, NUM_GENDER_CLASSES),
|
|
_one_hot(ethnicity_index, NUM_ETHNICITY_CLASSES),
|
|
],
|
|
axis=-1,
|
|
)
|
|
labels = np.repeat(labels, num_samples, axis=0)
|
|
z = rng.normal(size=(num_samples, self._decoder.latent_dim)).astype(
|
|
np.float32
|
|
)
|
|
return self._decoder.predict(z, labels)
|
|
|
|
|
|
class ExpressionSampler:
|
|
"""Numpy port of gnm.shape.semantic_sampler.ExpressionSampler."""
|
|
|
|
def __init__(self, decoder_path=EXPRESSION_DECODER_PATH):
|
|
self._decoder = _MLPDecoder(decoder_path)
|
|
|
|
def sample(self, expression_index, num_samples=1, rng=None):
|
|
rng = rng or np.random.default_rng()
|
|
labels = np.repeat(
|
|
_one_hot(expression_index, NUM_EXPRESSION_CLASSES), num_samples, axis=0
|
|
)
|
|
z = rng.normal(size=(num_samples, self._decoder.latent_dim)).astype(
|
|
np.float32
|
|
)
|
|
return self._decoder.predict(z, labels)
|