Implementing PAG (Progressive Attention Guidance) in Sana Inference

Progressive Attention Guidance (PAG) is a lightweight plug-in for classifier-free guidance that routes attention through a perturbed identity path, improving fine-detail fidelity in Sana diffusion models without sacrificing global coherence.

The NVlabs/Sana repository implements PAG as a drop-in augmentation to standard diffusion sampling. By introducing a separate perturbed token stream alongside unconditional and conditional inputs, PAG allows the model to progressively focus attention during inference. This guide covers the complete implementation path from configuration files to the attention processors in diffusion/model/nets/sana_blocks.py.

How PAG Works in Sana

PAG operates by injecting a third path—distinct from the standard conditional and unconditional routes—into the linear attention layers of the Sana transformer. This perturbed path receives an identity-attention treatment that preserves spatial structure while providing a separate gradient signal for guidance.

Core Architecture Components

The implementation centers on two specialized attention processors located in diffusion/model/nets/sana_blocks.py:

• PAGCFGIdentitySelfAttnProcessorLiteLA (lines 575-624): Handles the full three-stream case when guidance_type is set to "classifier-free_PAG" or "classifier-free_PAG_seq". It processes unconditional tokens (x_uncond), conditional tokens (x_org), and perturbed tokens (x_ptb) simultaneously.

• PAGIdentitySelfAttnProcessorLiteLA (lines 627-676): A leaner variant used when CFG is disabled but PAG remains active, processing only conditional and perturbed inputs.

The Perturbed Attention Mechanism

Inside these processors, input tensors are chunked into three parts. While the conditional path follows standard linear attention (qkv → kernel → attn_matmul → proj), the perturbed path bypasses the attention matrix entirely. Instead, it extracts only the V weight matrix from the QKV projection (v_weight = self.attn.qkv.weight[C*2:C*3]), applies optional bias, and routes directly to the output projection. This creates an identity attention effect that maintains spatial coherence while allowing the sampler to mix signals according to the pag_guidance_scale.

Implementation Steps

Using SanaPAGPipeline

The simplest entry point is the SanaPAGPipeline class, which subclasses the standard SanaPipeline and automatically configures the scheduler for PAG support.

from diffusers import SanaPAGPipeline
import torch

pipe = SanaPAGPipeline.from_pretrained(
    "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers",
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")

Configuring Guidance Parameters

After loading, set the guidance type and PAG scale through the scheduler config or pipeline arguments:

# Configure the DPM solver for PAG

pipe.scheduler.config.guidance_type = "classifier-free_PAG"  # or "classifier-free_PAG_seq"

pipe.scheduler.config.pag_guidance_scale = 2.5               # Typical range: 1.0-10.0

# Run inference with both CFG and PAG active

output = pipe(
    prompt="a futuristic cityscape at sunset",
    height=1024,
    width=1024,
    num_inference_steps=20,
    guidance_scale=4.5,          # Standard CFG scale

    pag_guidance_scale=2.5,      # PAG-specific scale

    generator=torch.Generator(device="cuda").manual_seed(42),
)

Selective Layer Application

PAG does not apply to every transformer block by default. The model configuration files (e.g., configs/sana_config/1024ms/Sana_1600M_img1024.yaml) contain the pag_applied_layers parameter that specifies which blocks receive the PAG processors. The dpm_solver.py (lines 486-558) routes these specifications during sampler initialization, attaching PAG processors only to the designated layers.

Code Examples

Full Inference Script

This complete example demonstrates proper PAG initialization with the sequential guidance variant:

import torch
from diffusers import SanaPAGPipeline

# Load PAG-enabled pipeline

pipe = SanaPAGPipeline.from_pretrained(
    "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers",
    torch_dtype=torch.bfloat16
).to("cuda")

# Set sequential PAG application for finer control

pipe.scheduler.config.guidance_type = "classifier-free_PAG_seq"
pipe.scheduler.config.pag_guidance_scale = 3.0

# Generate with PAG guidance

image = pipe(
    prompt="highly detailed cyberpunk street market, neon lights, raining",
    height=1024,
    width=1024,
    num_inference_steps=20,
    guidance_scale=4.5,
    pag_guidance_scale=3.0,
).images[0]

image.save("sana_pag_output.png")

Gradio UI Configuration

For interface development, reference the implementation in scripts/interface.py (lines 340-347):

import gradio as gr

with gr.Row():
    pag_scale = gr.Slider(
        label="PAG Guidance scale",
        minimum=1.0,
        maximum=10.0,
        value=2.5,
        step=0.5
    )
    guidance_type = gr.Dropdown(
        choices=["classifier-free", "classifier-free_PAG", "classifier-free_PAG_seq"],
        value="classifier-free_PAG_seq",
        label="Guidance type",
    )

Technical Deep Dive

Attention Processor Classes

The sana_blocks.py file implements the critical __call__ methods that handle tensor chunking. When PAGCFGIdentitySelfAttnProcessorLiteLA receives batched inputs, it splits tensors into three chunks, processes the perturbed chunk through the identity path (using only the V projection weights), then concatenates the results (torch.cat([out, x_ptb])) before returning to the sampler.

DPM Solver Integration

In diffusion/model/dpm_solver.py, the sampler checks guidance_type during the denoising loop. When the type is "classifier-free_PAG", it maintains three separate noise predictions (unconditional, conditional, and perturbed) and combines them using the formula:

eps_final = eps_uncond + guidance_scale * (eps_cond - eps_uncond) + pag_guidance_scale * (eps_cond - eps_ptb)

For "classifier-free_PAG_seq", the PAG guidance applies sequentially across the layers specified in pag_applied_layers, creating the progressive effect where early steps emphasize CFG and later steps increasingly weight the perturbed attention.

Summary

• PAG augments CFG by adding a perturbed identity-attention path to specific transformer blocks in Sana models.
• Two processors handle different modes: PAGCFGIdentitySelfAttnProcessorLiteLA for full CFG+PAG, and PAGIdentitySelfAttnProcessorLiteLA for PAG-only.
• Configuration happens at three levels: YAML configs select layers (pag_applied_layers), the scheduler selects the guidance type (classifier-free_PAG or classifier-free_PAG_seq), and the pipeline sets the intensity (pag_guidance_scale).
• Implementation requires the SanaPAGPipeline class and proper DPM solver configuration as shown in app/sana_pipeline.py.

Frequently Asked Questions

What is the difference between PAG and standard CFG in Sana?

Standard classifier-free guidance (CFG) computes the noise prediction using only unconditional and conditional inputs. PAG introduces a third perturbed input that bypasses the standard attention matrix, creating an identity-attention path. According to the source code in sana_blocks.py, this perturbed path uses only the V projection weights, providing a distinct signal that improves fine-detail generation without the artifacts that can occur with high CFG scales.

When should I use classifier-free_PAG versus classifier-free_PAG_seq?

Use classifier-free_PAG (parallel application) when you want PAG applied simultaneously across all designated layers in a single step, which offers faster inference. Use classifier-free_PAG_seq (sequential application) when you want PAG to apply progressively through the layer stack, as implemented in dpm_solver.py lines 486-558. The sequential mode typically yields better fine-grained detail but requires slightly more computation.

How do I tune the pag_guidance_scale parameter?

The pag_guidance_scale typically ranges from 1.0 to 10.0, with 2.5 serving as a common default. Lower values (1.0-2.0) provide subtle texture improvements, while higher values (4.0+) aggressively sharpen details but may introduce local inconsistencies. According to the Gradio interface in scripts/interface.py, values are usually adjusted in 0.5 increments alongside the standard guidance_scale (typically 4.5 for CFG).

Can I apply PAG to specific layers only?

Yes. The pag_applied_layers parameter in the model configuration YAML files (e.g., Sana_1600M_img1024.yaml line 39) controls which transformer blocks receive PAG processors. The DPM solver reads this list during initialization and attaches PAGCFGIdentitySelfAttnProcessorLiteLA only to the specified layer indices, leaving other layers to use standard attention. This selective application is crucial for managing computational overhead while maximizing visual quality.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}