azula.plugins.flux

Flux plugin.

This plugin depends on diffusers and transformers. To use it, install the dependencies in your environment

pip install diffusers transformers accelerate protobuf sentencepiece

before importing the plugin.

from azula.plugins import flux

References

FLUX (Black Forest Labs, 2024)

https://github.com/black-forest-labs/flux

Classes

AutoEncoder	Creates an auto-encoder wrapper.
TextEncoder	Creates a text encoder.
FluxDenoiser	Creates a Flux denoiser.

Functions

load_model

Loads a pre-trained Flux latent denoiser.

Descriptions

class azula.plugins.flux.AutoEncoder(vae, shift=0.0, scale=1.0)

Creates an auto-encoder wrapper.

encode(x)

Encodes images to latents.

Parameters:

x (Tensor) – A batch of images \(x\), with shape \((B, 3, H, W)\). Pixel values are expected to range between -1 and 1.

Returns:

A batch of latents \(z \sim q(Z \mid x)\), with shape \((B, H / 16, W / 16, 64)\).

Return type:

Tensor

decode(z)

Decodes latents to images.

Parameters:

z (Tensor) – A batch of latents \(z\), with shape \((B, H / 16, W / 16, 64)\).

Returns:

A batch of images \(x = D(z)\), with shape \((B, 3, H, W)\).

Return type:

Tensor

class azula.plugins.flux.TextEncoder(clip, clip_tokenizer, t5, t5_tokenizer)

Creates a text encoder.

forward(prompt)

Parameters:

prompt (str | Sequence[str]) – A text prompt or list of text prompts.

Returns:

The CLIP and T5 encoded prompt(s).

Return type:

dict[str, Tensor]

class azula.plugins.flux.FluxDenoiser(backbone, schedule=None)

Creates a Flux denoiser.

Parameters:

• backbone (Module) – A time conditional network.

• schedule (Schedule) – A noise schedule. If None, use azula.noise.DecaySchedule instead.

forward(z_t, t, prompt_clip, prompt_t5, guidance=4.0, **kwargs)

Parameters:

• z_t (Tensor) – A noisy tensor \(z_t\), with shape \((B, H, W, 64)\).

• t (Tensor) – The time \(t\), with shape \(()\) or \((B)\).

• prompt_clip (Tensor) – The CLIP-encoded text prompt \(y\), with shape \((B, F)\).

• prompt_t5 (Tensor) – The T5-encoded text prompt \(y\), with shape \((B, L, D)\).

• guidance (float | Tensor) – The guidance strength \(\omega \in \mathbb{R}\).

• kwargs – Optional keyword arguments.

Returns:

The Dirac delta \(\delta(Z - \mu_\phi(z_t \mid y))\).

Return type:

DiracPosterior

azula.plugins.flux.load_model(name='flux_1_dev', **kwargs)

Loads a pre-trained Flux latent denoiser.

Parameters:

• name (str) – The pre-trained model name.

• kwargs – Keyword arguments passed to diffusers.FluxPipeline.from_pretrained.

Returns:

A pre-trained latent denoiser and the corresponding auto-encoder and text encoder.

Return type:

tuple[Denoiser, AutoEncoder, TextEncoder]