Files
NobodyandClaude Fable 5 6e4098f51e Initial release: Google GNM Head v3 as ComfyUI custom nodes
- 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>
2026-07-15 19:46:46 +09:00

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)