gensbi.models.flux1#

Submodules#

Classes#

Flux

Transformer model for flow matching on sequences.

FluxCFMLoss

FluxCFMLoss is a class that computes the continuous flow matching loss for the Flux model.

FluxDiffLoss

FluxDiffLoss is a class that computes the diffusion score matching loss for the Flux model.

FluxParams

Parameters for the Flux model.

FluxWrapper

This class is used to wrap around another model. We define a call method which returns the model output.

Package Contents#

class gensbi.models.flux1.Flux(params)[source]#

Bases: flax.nnx.Module

Transformer model for flow matching on sequences.

Parameters:

params (FluxParams)

__call__(t, obs, obs_ids, cond, cond_ids, conditioned=True, guidance=None)[source]#
Parameters:

  • t (jax.Array)

  • obs (jax.Array)

  • obs_ids (jax.Array)

  • cond (jax.Array)

  • cond_ids (jax.Array)

  • conditioned (bool | jax.Array)

  • guidance (jax.Array | None)

Return type:

jax.Array

cond_in#
condition_embedding#
condition_null#
double_blocks#
final_layer#
hidden_size#
in_channels#
num_heads#
obs_in#
out_channels#
params#
pe_embedder#
qkv_features#
single_blocks#
time_in#
vector_in#
class gensbi.models.flux1.FluxCFMLoss(path, reduction='mean', cfg_scale=None)[source]#

Bases: gensbi.flow_matching.loss.ContinuousFMLoss

FluxCFMLoss is a class that computes the continuous flow matching loss for the Flux model.

Parameters:

  • path – Probability path (x-prediction training).

  • reduction (str, optional) – Specify the reduction to apply to the output 'none' | 'mean' | 'sum'. 'none': no reduction is applied to the output, 'mean': the output is reduced by mean over sequence elements, 'sum': the output is reduced by sum over sequence elements. Defaults to ‘mean’.

__call__(vf, batch, cond, obs_ids, cond_ids)[source]#

Evaluates the continuous flow matching loss.

Parameters:

  • vf (callable) – The vector field model to evaluate.

  • batch (tuple) – A tuple containing the input data (x_0, x_1, t).

  • cond (jnp.ndarray) – The conditioning data.

  • obs_ids (jnp.ndarray) – The observation IDs.

  • cond_ids (jnp.ndarray) – The conditioning IDs.

Returns:

The computed loss.

Return type:

jnp.ndarray

cfg_scale = None#
class gensbi.models.flux1.FluxDiffLoss(path)[source]#

Bases: flax.nnx.Module

FluxDiffLoss is a class that computes the diffusion score matching loss for the Flux model.

Parameters:

path – Probability path for training.

__call__(key, model, batch, cond, obs_ids, cond_ids)[source]#

Evaluate the continuous flow matching loss.

Parameters:

  • key (jax.random.PRNGKey) – Random key for stochastic operations.

  • model (Callable) – F model.

  • batch (Tuple[Array, Array, Array]) – Input data (x_1, sigma).

  • cond (jnp.ndarray) – The conditioning data.

  • obs_ids (jnp.ndarray) – The observation IDs.

  • cond_ids (jnp.ndarray) – The conditioning IDs.

Returns:

Computed loss.

Return type:

Array

loss_fn#
path#
class gensbi.models.flux1.FluxParams[source]#

Parameters for the Flux model.

Parameters:

  • in_channels (int) – Number of input channels.

  • vec_in_dim (Union[int, None]) – Dimension of the vector input, if applicable.

  • context_in_dim (int) – Dimension of the context input.

  • mlp_ratio (float) – Ratio for the MLP layers.

  • num_heads (int) – Number of attention heads.

  • depth (int) – Number of double stream blocks.

  • depth_single_blocks (int) – Number of single stream blocks.

  • axes_dim (list[int]) – Dimensions of the axes for positional encoding.

  • qkv_bias (bool) – Whether to use bias in QKV layers.

  • rngs (nnx.Rngs) – Random number generators for initialization.

  • obs_dim (int) – Observation dimension.

  • cond_dim (int) – Condition dimension.

  • theta (int) – Scaling factor for positional encoding.

  • guidance_embed (bool) – Whether to use guidance embedding.

  • qkv_multiplier (int) – Multiplier for QKV features.

  • param_dtype (DTypeLike) – Data type for model parameters.

__post_init__()[source]#
axes_dim: list[int]#
cond_dim: int#
context_in_dim: int#
depth: int#
depth_single_blocks: int#
guidance_embed: bool = False#
in_channels: int#
mlp_ratio: float#
num_heads: int#
obs_dim: int#
param_dtype: jax.typing.DTypeLike#
qkv_bias: bool#
qkv_multiplier: int = 1#
rngs: flax.nnx.Rngs#
theta: int = 10000#
vec_in_dim: int | None#
class gensbi.models.flux1.FluxWrapper(model)[source]#

Bases: gensbi.utils.model_wrapping.ModelWrapper

This class is used to wrap around another model. We define a call method which returns the model output. Furthermore, we define a vector_field method which computes the vector field of the model, and a divergence method which computes the divergence of the model, in a form useful for diffrax. This is useful for ODE solvers that require the vector field and divergence of the model.

__call__(t, obs, obs_ids, cond, cond_ids, conditioned=True, guidance=None)[source]#

This method defines how inputs should be passed through the wrapped model. Here, we’re assuming that the wrapped model takes both \(obs\) and \(t\) as input, along with any additional keyword arguments.

Optional things to do here:

  • check that t is in the dimensions that the model is expecting.

  • add a custom forward pass logic.

  • call the wrapped model.

given obs, t
returns the model output for input obs at time t, with extra information extra.
Parameters:

  • obs (Array) – input data to the model (batch_size, …).

  • t (Array) – time (batch_size).

  • **extras – additional information forwarded to the model, e.g., text condition.

  • obs_ids (jax.Array)

  • cond (jax.Array)

  • cond_ids (jax.Array)

  • conditioned (bool | jax.Array)

  • guidance (jax.Array | None)

Returns:

model output.

Return type:

Array