class
StableDiffusionForImageGeneration
extends
ImageGenerationModelStableDiffusionForImageGeneration(config: StableDiffusionConfig)Stable Diffusion posed as a diffusion model: the objective and a sampler.
Parameters
configStableDiffusionConfigThe variant to build.
Attributes
stable_diffusionStableDiffusionModelThe autoencoder, U-Net and noising schedule.
Notes
Reference: Rombach et al., CVPR 2022 (arXiv:2112.10752), §4.3.
forward is one training step of the denoising objective, as
in every other image-generation family; generate runs the
reverse process and decodes. Conditioning arrives already encoded,
and for classifier-free guidance the caller supplies the
unconditional sequence too — usually the text encoder's output for
the empty string, which is not the same thing as zeros.
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... StableDiffusionConfig, StableDiffusionForImageGeneration)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... attention_head_dim=32,
... cross_attention_dim=16, context_length=4)
>>> model = StableDiffusionForImageGeneration(config).eval()
>>> context = lucid.randn((1, 4, 16))
>>> model.generate(context, num_inference_steps=2).samples.shape
(1, 3, 32, 32)
Calling the model trains it — the loss is always reported:
>>> out = model(lucid.randn((1, 3, 32, 32)), context)
>>> out.loss.shape, out.noise_pred.shape
((), (1, 4, 8, 8))Used by 2
Constructors
1Instance methods
2forward(images: Tensor, context: Tensor, timestep: Tensor | None = None)One training step of the denoising objective.
Parameters
Returns
StableDiffusionOutputThe prediction, the noised latent, and the loss — always present, since a task model is called to be trained.
generate(context: Tensor, uncond_context: Tensor | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, eta: float = 0.0, latent: Tensor | None = None, sampler: str = 'pndm')Sample an image from the conditioning.
Parameters
contextTensor(B, L, cross_attention_dim).The unconditional sequence. Guidance is skipped when absent,
which is equivalent to
guidance_scale = 1.num_inference_stepsint= 50Network evaluations along the reverse trajectory.
guidance_scalefloat= 7.5 in
.
etafloat= 0.0DDIM at 0, DDPM at 1. Ignored by PNDM, which is
deterministic by construction.
Starting noise. Drawn when absent.
samplerstr= "pndm""pndm" reproduces the released pipeline — its
model_index.json names PNDMScheduler. "ddim" is
the paper's, and the one eta applies to.Returns
GenerationOutputsamples of shape (B, out_channels, sample_size, sample_size) — decoded pixels, not the latent.
Raises
ValueErrorIf
guidance_scale is negative, or the two conditioning
sequences disagree in shape.Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... StableDiffusionConfig, StableDiffusionForImageGeneration)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... attention_head_dim=32,
... cross_attention_dim=16, context_length=4)
>>> model = StableDiffusionForImageGeneration(config).eval()
>>> context = lucid.randn((1, 4, 16))
>>> start = lucid.randn((1, 4, 8, 8))
>>> image = model.generate(
... context, num_inference_steps=2, latent=start).samples
>>> image.shape
(1, 3, 32, 32)
PNDM draws no noise of its own, so the starting latent is the whole
seed: pass the same one and the same image comes back.
>>> again = model.generate(
... context, num_inference_steps=2, latent=start).samples
>>> bool((again == image).all())
True
At guidance_scale=1 the extrapolation is the conditional
prediction itself, so an unconditional sequence changes nothing; at
the default 7.5 it does.
>>> uncond = lucid.randn((1, 4, 16)) # stands in for the empty caption
>>> unguided = model.generate(context, uncond, num_inference_steps=2,
... guidance_scale=1.0, latent=start).samples
>>> bool((unguided == image).all())
True
>>> guided = model.generate(context, uncond, num_inference_steps=2,
... latent=start).samples
>>> bool((guided == image).all())
False