gensbi.models.flux1.loss

Classes

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.

Module Contents

class gensbi.models.flux1.loss.FluxCFMLoss(path, reduction='mean', cfg_scale=None)

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)

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.loss.FluxDiffLoss(path)

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)

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