"""Conditioning seams for the transformer flow.
Adapted from apple/ml-tarflow (TarFlow); see models/tarflow/LICENSE.apple.
Prefix-concatenation conditioning adapted from apple/ml-starflow (STARFlow); see models/tarflow/LICENSE.starflow.
``AdditiveBiasConditioner`` is the continuous analog of TarFlow's ``class_embed``:
an MLP embeds the condition to a ``channels``-vector that is broadcast-added to
every token. The signal depends only on the condition (constant w.r.t. the
modeled variable), so it shifts the affine params without breaking the
triangular Jacobian. A plain 2-layer MLP is used (not ``MLPEmbedder``, whose
``hidden_dim % in_dim == 0`` constraint does not fit arbitrary ``cond_dim``).
"""
import jax
import jax.numpy as jnp
from flax import nnx
from jax import Array
from jax.typing import DTypeLike
from gensbi.models.core.patching import patchify_2d
[docs]
class AdditiveBiasConditioner(nnx.Module):
"""Embed a vector condition as a per-token additive bias.
A two-layer MLP maps the condition to a ``channels``-dimensional vector
that is broadcast-added to every token in the sequence. When
``cond_dim == 0`` the conditioner is unconditional and :meth:`embed`
returns ``(None, None)``.
Parameters
----------
cond_dim : int
Condition dimensionality. Set to ``0`` for an unconditional model.
channels : int
Output channel width matching the transformer embedding dimension.
rngs : nnx.Rngs
Flax RNG container for linear layer initialization.
cond_channels : int, optional
Number of channels in the conditioning input ``(B, cond_dim, C_cond)``.
Default is ``1``. The input linear layer is widened to accept
``cond_dim * cond_channels`` features, folding the channel axis before
the MLP (same flattening performed in :meth:`embed`).
param_dtype : DTypeLike, optional
Dtype for the stored (master) kernel/bias parameters. Default is
``float32``.
dtype : DTypeLike, optional
Compute dtype forwarded to the MLP layers. Default is ``float32``,
matching ``param_dtype``, so with default arguments this is a
bit-identical no-op cast.
"""
def __init__(self, cond_dim: int, channels: int, rngs: nnx.Rngs,
cond_channels: int = 1, param_dtype: DTypeLike = jnp.float32,
dtype: DTypeLike = jnp.float32):
[docs]
self.cond_dim = cond_dim
[docs]
self.cond_channels = cond_channels
if cond_dim > 0:
self.l1 = nnx.Linear(cond_dim * cond_channels, channels, rngs=rngs,
param_dtype=param_dtype, dtype=dtype)
self.l2 = nnx.Linear(channels, channels, rngs=rngs,
param_dtype=param_dtype, dtype=dtype)
[docs]
def embed(self, cond: Array | None):
"""Embed the condition into a per-token additive bias.
Parameters
----------
cond : Array or None
Condition vector of shape ``(B, cond_dim)`` or
``(B, cond_dim, C_cond)``, or ``None`` when the model is
unconditional (``cond_dim == 0``).
Returns
-------
bias : Array or None
Per-token additive bias of shape ``(B, channels)``, or ``None``
when ``cond_dim == 0``.
prefix : None
This conditioner does not produce prefix tokens; always ``None``.
Raises
------
ValueError
If ``cond`` is ``None`` when ``cond_dim > 0``.
"""
if self.cond_dim == 0:
return (None, None)
if cond is None:
raise ValueError(
"cond is required: this conditioner was built with cond_dim > 0")
cond = jnp.asarray(cond).reshape(cond.shape[0], -1) # (B, cond_dim*C_cond)
bias = self.l2(jax.nn.silu(self.l1(cond)))
return (bias, None)
[docs]
class VectorConditioner(nnx.Module):
"""Embed a vector condition as one prefix token per coordinate.
Each of the ``cond_dim`` coordinates is a token of ``C_cond`` channels; a
shared ``Linear(cond_channels, channels)`` projects each to the transformer
width, plus per-coordinate positional embeddings. Produces ``M = cond_dim``
prefix tokens (no flatten).
Parameters
----------
cond_dim : int
Condition dimensionality (number of coordinates / prefix tokens).
cond_channels : int
Number of channels per coordinate in the input condition
``(B, cond_dim, cond_channels)``.
channels : int
Output channel width matching the transformer embedding dimension.
rngs : nnx.Rngs
Flax RNG container for linear layer and positional embedding
initialization.
param_dtype : DTypeLike, optional
Dtype for the stored (master) kernel/bias/positional parameters.
Default is ``float32``.
dtype : DTypeLike, optional
Compute dtype forwarded to ``proj``. Default is ``float32``,
matching ``param_dtype``, so with default arguments this is a
bit-identical no-op cast.
"""
def __init__(self, cond_dim: int, cond_channels: int, channels: int,
rngs: nnx.Rngs, param_dtype: DTypeLike = jnp.float32,
dtype: DTypeLike = jnp.float32):
[docs]
self.cond_dim = cond_dim
[docs]
self.cond_channels = cond_channels
[docs]
self.channels = channels
[docs]
self.proj = nnx.Linear(cond_channels, channels, rngs=rngs,
param_dtype=param_dtype, dtype=dtype)
[docs]
self.pos = nnx.Param(
(jax.random.normal(rngs.params(), (cond_dim, channels)) * 1e-2
).astype(param_dtype))
[docs]
def embed(self, cond: Array | None):
"""Embed the condition into per-coordinate prefix tokens.
Parameters
----------
cond : Array or None
Condition array of shape ``(B, cond_dim, C_cond)``.
Returns
-------
bias : None
This conditioner does not produce a per-token additive bias;
always ``None``.
prefix : Array
Prefix token sequence of shape ``(B, cond_dim, channels)`` with
learned positional embeddings added.
Raises
------
ValueError
If ``cond`` is ``None``.
"""
if cond is None:
raise ValueError("cond is required for VectorConditioner")
cond = jnp.asarray(cond) # (B, cond_dim, C_cond)
proj = self.proj(cond)
return (None, proj + self.pos[...].astype(proj.dtype)[None])
[docs]
class ImageConditioner(nnx.Module):
"""Embed an image condition as prefix tokens prepended to the sequence.
Patchifies a spatial image ``(B, H, W, C)`` into
``M = (H / patch_size) * (W / patch_size)`` flat patch vectors, projects
each patch to ``channels`` dimensions, and adds learned positional
embeddings.
Parameters
----------
cond_channels : int
Number of channels in the conditioning image.
patch_size : int
Spatial size of each square patch (height and width in pixels).
channels : int
Output channel width matching the transformer embedding dimension.
num_tokens : int
Number of prefix tokens; must equal
``(H / patch_size) * (W / patch_size)``.
rngs : nnx.Rngs
Flax RNG container for projection layer and positional embedding
initialization.
param_dtype : DTypeLike, optional
Dtype for the stored (master) kernel/bias/positional parameters.
Default is ``float32``.
dtype : DTypeLike, optional
Compute dtype forwarded to ``proj``. Default is ``float32``,
matching ``param_dtype``, so with default arguments this is a
bit-identical no-op cast.
"""
def __init__(self, cond_channels: int, patch_size: int, channels: int,
num_tokens: int, rngs: nnx.Rngs,
param_dtype: DTypeLike = jnp.float32,
dtype: DTypeLike = jnp.float32):
[docs]
self.patch_size = patch_size
[docs]
self.channels = channels
in_f = cond_channels * patch_size * patch_size
[docs]
self.proj = nnx.Linear(in_f, channels, rngs=rngs,
param_dtype=param_dtype, dtype=dtype)
[docs]
self.pos = nnx.Param(
(jax.random.normal(rngs.params(), (num_tokens, channels)) * 1e-2
).astype(param_dtype))
[docs]
def embed(self, cond: Array | None):
"""Patchify an image condition and embed it as prefix tokens.
Parameters
----------
cond : Array or None
Image condition of shape ``(B, H, W, C)``.
Returns
-------
bias : None
This conditioner does not produce a per-token additive bias;
always ``None``.
prefix : Array
Prefix token sequence of shape ``(B, num_tokens, channels)`` with
learned positional embeddings added.
Raises
------
ValueError
If ``cond`` is ``None``.
"""
if cond is None:
raise ValueError("cond is required for ImageConditioner")
patches = patchify_2d(cond, size=self.patch_size) # (B, M, in_f)
proj = self.proj(patches)
return (None, proj + self.pos[...].astype(proj.dtype)[None])
