diffusion

diffusion

DiffusionBatcher

DiffusionBatcher(envoy: Any, kwargs: Optional[dict] = None)

Bases: Batcher

Batcher for the denoiser's guidance-doubled, multi-image batch layout.

Where the base batcher assumes every batched activation is a plain dim-0 stack of the invokes' rows, a diffusion denoiser sees an expanded batch: each prompt is repeated num_images_per_prompt times, and under classifier-free guidance the whole thing is doubled (an unconditional half followed by a conditional half). This maps each invoke's plain [start, size] group onto that layout, picking the case by the tensor's leading dim at run time — so an intervention on model.unet reads (and writes) exactly its invoke's rows across both halves.

num_images instance-attribute

num_images: int = self.kwargs.get('num_images_per_prompt', 1)

image_groups instance-attribute

image_groups: dict[int, list] = {}

image_total instance-attribute

image_total = 0

add

add(*inputs: Any, **kwargs: Any) -> BatchGroup

DiffusionModel

DiffusionModel(repo_id: Any, *args: Any, automodel: Any = None, **kwargs: Any)

Bases: HuggingFaceModel

A model backed by a diffusers.DiffusionPipeline.

A diffusion pipeline orchestrates several modules (unet/transformer, vae, text_encoder, ...) around a denoising loop. This wraps the whole pipeline as one nnsight model: the pipeline is exposed as pipeline, each of its module components as an envoy (model.unet, model.vae, ...), and interventions apply to any component along the way.

Both trace and generate run the whole pipeline, with interventions firing on every component the denoising loop invokes; model.output (and tracer.result) is the pipeline's output object.

A traced run defaults to num_inference_steps=1 — a fast one-step pass for inspecting or editing activations — and with model.generate(...): is a traced run, so it takes that default too; only a bare model.generate(...) call outside a trace uses the pipeline's own default. The same call means one denoising step inside a with and fifty outside it, so pass num_inference_steps= whenever the image itself matters. To run one component's forward on its own, trace that envoy directly — with model.unet.trace(sample, timestep, encoder_hidden_states=...):.

On dispatch a real pipeline is built with real weights. The lazy meta build can't load a pipeline without weights, so each module component is constructed from its config on the meta device while the light components (scheduler, tokenizer, ...) load normally, and a meta pipeline of the same shape is assembled from them.

Requires the optional diffusers package.

PARAMETER	DESCRIPTION
repo_id	A diffusers pipeline repo id (or local path) to load. TYPE: Any
*args	Forwarded to the mixin chain. TYPE: Any DEFAULT: ()
**kwargs	Forwarded to the mixin chain and, on dispatch, to DiffusionPipeline.from_pretrained. rename= exposes components under different envoy names (e.g. {"unet": "denoiser"}); dispatch=True loads real weights immediately instead of building lazily on the meta device. TYPE: Any DEFAULT: {}

Examples:

>>> from nnsight.modeling.diffusion import DiffusionModel
>>> model = DiffusionModel("hf-internal-testing/tiny-stable-diffusion-torch")
>>> with model.generate("a photo of a cat", num_inference_steps=2) as tracer:
...     latents = model.unet.output[0].save()   # per denoising step
...     images = model.output.save()
>>> images.images[0]  # a PIL image

pipeline instance-attribute

pipeline: Optional['DiffusionPipeline'] = None

automodel instance-attribute

automodel = automodel

trace

trace(*inputs: Any, **kwargs: Any)

Trace the whole pipeline, defaulting to a single denoising step.

num_inference_steps=1 unless overridden — a fast one-step pass for inspecting or editing activations. model.output is the pipeline's output object. with model.generate(...): routes here as well, so a traced generation takes the one-step default too.

Examples:

>>> with model.trace("a photo of a cat"):
...     latents = model.unet.output[0].save()
...     images = model.output.save()

generate

generate(*inputs: Any, **kwargs: Any) -> Any

Run the diffusion pipeline, returning its output object.

with model.generate(...): traces the whole pipeline, so the block's interventions run against every component the denoising loop invokes (use tracer.iter to target a particular inference step); calling it directly just runs the pipeline. A traced run goes through trace, so it defaults to one denoising step where a direct call takes the pipeline's own default — name num_inference_steps to fix the count either way. The return value is the pipeline's own output object — read the images off its .images (or off model.output / tracer.result inside a trace).

Examples:

>>> with model.generate("a photo of a cat", num_inference_steps=20):
...     images = model.output.save()
>>> images.images[0]  # a PIL image

PARAMETER	DESCRIPTION
*inputs	The pipeline's inputs, e.g. a text prompt. TYPE: Any DEFAULT: ()
**kwargs	Forwarded to the pipeline, e.g. num_inference_steps, guidance_scale, num_images_per_prompt, output_type. seed (int) is turned into a reproducible generator — a per-image list for a batch; pass generator directly to override. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
Any	The pipeline's output object; its .images holds the generated images.

__getstate__

__getstate__() -> dict