pag model/pipeline review

`pag` model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules. Reviewed PAG public exports/lazy imports, pipeline/runtime behavior, related base-pipeline consistency, docs/examples, and tests/pipelines/pag.

Issue 1: `pag_applied_layers="blocks.1"` also matches `blocks.10`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pag_utils.py#L58-L77

Problem:
re.search(layer_id, name) is used directly, and the numeric disambiguation only compares the last dot-separated tokens. For names like blocks.10.attn1, the last token is attn1, so blocks.1 still matches blocks.10. This contradicts the nearby comment and over-applies PAG to unintended transformer blocks.

Impact:
Users selecting a single DiT block can silently perturb additional blocks, changing quality/performance and making layer ablations unreliable.

Reproduction:

import torch.nn as nn
from diffusers.models.attention_processor import Attention
from diffusers.pipelines.pag.pag_utils import PAGMixin

class TinyTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([nn.Module() for _ in range(11)])
        for block in self.blocks:
            block.attn1 = Attention(query_dim=8, heads=1, dim_head=8)

    @property
    def attn_processors(self):
        return {f"{n}.processor": m.processor for n, m in self.named_modules() if isinstance(m, Attention)}

class Dummy(PAGMixin):
    def __init__(self):
        self.transformer = TinyTransformer()
        self.set_pag_applied_layers(["blocks.1"])

pipe = Dummy()
pipe._set_pag_attn_processor(pipe.pag_applied_layers, do_classifier_free_guidance=False)
print(sorted(pipe.pag_attn_processors))
# Includes both blocks.1.attn1.processor and blocks.10.attn1.processor

Relevant precedent:
No duplicate found in GitHub issue/PR searches for PAG pag_applied_layers blocks.1 blocks.10.

Suggested fix:

match = re.search(layer_id, name)
if match is None:
    continue
if layer_id[-1].isdigit() and match.end() < len(name) and name[match.end()].isdigit():
    continue
if is_self_attn(module):
    target_modules.append(module)

Issue 2: PAG processors are not restored if generation raises

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sd.py#L996-L1001
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sd.py#L1073-L1077

Problem:
PAG pipelines save the original attention processors before the denoising loop and restore them only on the normal success path. If a callback, scheduler, VAE, or user interrupt raises after _set_pag_attn_processor, the model keeps PAG processors installed.

Impact:
A later pag_scale=0 call can fail or produce wrong results because the UNet/transformer still expects PAG-expanded batches.

Reproduction:

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPAGPipeline, UNet2DConditionModel

unet = UNet2DConditionModel(block_out_channels=(4, 8), layers_per_block=2, sample_size=32, in_channels=4, out_channels=4,
    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
    cross_attention_dim=8, norm_num_groups=2)
vae = AutoencoderKL(block_out_channels=[4, 8], in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D"]*2,
    up_block_types=["UpDecoderBlock2D"]*2, latent_channels=4, norm_num_groups=2)
text_encoder = CLIPTextModel(CLIPTextConfig(bos_token_id=0, eos_token_id=2, hidden_size=8, intermediate_size=16,
    num_attention_heads=2, num_hidden_layers=2, pad_token_id=1, vocab_size=1000))
pipe = StableDiffusionPAGPipeline(
    unet=unet, scheduler=DDIMScheduler(), vae=vae, text_encoder=text_encoder,
    tokenizer=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"),
    safety_checker=None, feature_extractor=None,
)
pipe.set_progress_bar_config(disable=True)
try:
    pipe("x", num_inference_steps=1, guidance_scale=5, pag_scale=1, output_type="latent",
         callback_on_step_end=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")))
except RuntimeError:
    pass
pipe("x", num_inference_steps=1, guidance_scale=5, pag_scale=0, output_type="latent")
# ValueError: not enough values to unpack (expected 3, got 2)

Relevant precedent:
No duplicate found for PAG callback exception attention processor restore.

Suggested fix:

original_attn_proc = None
try:
    if self.do_perturbed_attention_guidance:
        original_attn_proc = self.unet.attn_processors
        self._set_pag_attn_processor(self.pag_applied_layers, self.do_classifier_free_guidance)
    # denoise/decode body
finally:
    if original_attn_proc is not None:
        self.unet.set_attn_processor(original_attn_proc)

Issue 3: SD3 PAG pipelines are stale versus base SD3

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py#L136-L176
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py#L686-L715
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py#L152-L191
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py#L749-L768

Problem:
Base SD3 supports mu for dynamic shifting, SD3 IP-Adapter loading/call inputs, and text-to-image skip-layer guidance. The PAG SD3 variants do not expose those APIs. They also expose callback tensor inputs for negative pooled embeds while omitting the consumed pooled_prompt_embeds.

Impact:
SD3.5-style schedulers with use_dynamic_shifting=True fail because PAG cannot pass mu; IP-Adapter workflows cannot be used with SD3 PAG; callbacks cannot modify pooled conditioning.

Reproduction:

import inspect
from diffusers import FlowMatchEulerDiscreteScheduler, StableDiffusion3PAGPipeline, StableDiffusion3PAGImg2ImgPipeline

for cls in (StableDiffusion3PAGPipeline, StableDiffusion3PAGImg2ImgPipeline):
    params = inspect.signature(cls.__call__).parameters
    print(cls.__name__, "mu" in params, "ip_adapter_image" in params, cls._callback_tensor_inputs)

scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
scheduler.set_timesteps(2)
# ValueError: `mu` must be passed when `use_dynamic_shifting` is set to be `True`

Relevant precedent:
Base SD3 implements these APIs:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L794-L807
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L1012-L1037

No duplicate found for StableDiffusion3PAGPipeline mu use_dynamic_shifting or SD3 PAG IP-Adapter searches.

Suggested fix:
Port the current base SD3 __call__ API and logic into both SD3 PAG variants, then add the PAG batch expansion around the updated prompt/pooled/IP-Adapter conditioning. Include SD3IPAdapterMixin, mu, and the base callback tensor contract.

Issue 4: ControlNet PAG variants dropped `guess_mode`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py#L980-L1006
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py#L1235-L1260
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py#L1003-L1044
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py#L1512-L1519

Problem:
StableDiffusionControlNetPAGInpaintPipeline and StableDiffusionXLControlNetPAGPipeline omit the public guess_mode argument and hard-code guess_mode=False in image preparation and ControlNet forward. They also skip the base pipeline’s global_pool_conditions fallback.

Impact:
ControlNet checkpoints that require guess mode/global pooling cannot be used through these PAG variants, and users cannot request a feature supported by the corresponding base pipelines.

Reproduction:

import inspect
from diffusers import StableDiffusionControlNetPAGInpaintPipeline, StableDiffusionXLControlNetPAGPipeline

for cls in (StableDiffusionControlNetPAGInpaintPipeline, StableDiffusionXLControlNetPAGPipeline):
    print(cls.__name__, "guess_mode" in inspect.signature(cls.__call__).parameters)
# False for both

Relevant precedent:
Base inpaint and SDXL ControlNet expose and normalize guess_mode:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py#L1020-L1027
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py#L1025-L1032

Search found old base SDXL guess-mode issue #4709, but it is closed and not a duplicate of these PAG omissions.

Suggested fix:

# __call__ signature
guess_mode: bool = False,

# after resolving controlnet
global_pool_conditions = (
    controlnet.config.global_pool_conditions
    if isinstance(controlnet, ControlNetModel)
    else controlnet.nets[0].config.global_pool_conditions
)
guess_mode = guess_mode or global_pool_conditions

# replace every hard-coded guess_mode=False with guess_mode=guess_mode

Issue 5: `SanaPAGPipeline` lost Sana LoRA and attention kwargs support

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sana.py#L148-L160
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sana.py#L650-L687
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pag/pipeline_pag_sana.py#L906-L910

Problem:
Base SanaPipeline inherits SanaLoraLoaderMixin, accepts attention_kwargs, uses attention_kwargs["scale"] for prompt encoding, and forwards the kwargs into the transformer. SanaPAGPipeline inherits only DiffusionPipeline, PAGMixin and has no attention_kwargs argument or forwarding.

Impact:
Users cannot load or scale Sana LoRAs with the PAG pipeline, despite Sana supporting them and the PAG API page carrying a LoRA badge.

Reproduction:

import inspect
from diffusers import SanaPipeline, SanaPAGPipeline

print(hasattr(SanaPipeline, "load_lora_weights"), hasattr(SanaPAGPipeline, "load_lora_weights"))
print("attention_kwargs" in inspect.signature(SanaPipeline.__call__).parameters)
print("attention_kwargs" in inspect.signature(SanaPAGPipeline.__call__).parameters)
# True False
# True
# False

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana.py#L190-L190
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana.py#L729-L753
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana.py#L888-L903

Related SanaPAG quality issue #10241 exists, but it is not a duplicate of the LoRA/attention kwargs API gap.

Suggested fix:

class SanaPAGPipeline(DiffusionPipeline, SanaLoraLoaderMixin, PAGMixin):
    ...

def __call__(..., attention_kwargs: dict[str, Any] | None = None, ...):
    self._attention_kwargs = attention_kwargs
    lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None
    ...
    noise_pred = self.transformer(..., attention_kwargs=self.attention_kwargs, return_dict=False)[0]

Issue 6: PAG docs have broken examples and omit Sana API docs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/using-diffusers/pag.md#L124-L138
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/using-diffusers/pag.md#L167-L193
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/api/pipelines/pag.md#L36-L113

Problem:
The guide defines pag_scales/guidance_scales but calls pag_scale/guidance_scale, references AutoPipelineForInpaiting, and uses pipeline_t2i where the preceding snippet defines pipeline_pag. The API page exports SanaPAGPipeline in code but never documents it.

Impact:
Copy-pasted PAG guide snippets fail before generation, and Sana PAG users cannot find the API reference.

Reproduction:

from pathlib import Path

guide = Path("docs/source/en/using-diffusers/pag.md").read_text()
api = Path("docs/source/en/api/pipelines/pag.md").read_text()
assert "AutoPipelineForInpaiting" not in guide
assert "pag_scales" not in guide or "pag_scale=pag_scale" not in guide
assert "## SanaPAGPipeline" in api

Relevant precedent:
No duplicate found for AutoPipelineForInpaiting pag_scale guidance_scale PAG docs.

Suggested fix:

pag_scale = 4.0
guidance_scale = 7.0
...
pipeline = AutoPipelineForInpainting.from_pipe(pipeline_t2i, enable_pag=True)

Also add:

## SanaPAGPipeline
[[autodoc]] SanaPAGPipeline
  - all
  - __call__

Issue 7: Slow coverage is missing for most PAG variants

Problem:
Fast tests exist across the family, but these PAG variants have no @slow integration tests. The prompt explicitly requires missing slow tests to be reported.

Impact:
Real checkpoint/API drift is not covered for many PAG pipelines, including the stale SD3/Sana/ControlNet gaps above.

Reproduction:

from pathlib import Path

missing = [
    p.as_posix()
    for p in sorted(Path("tests/pipelines/pag").glob("test_pag_*.py"))
    if "@slow" not in p.read_text(encoding="utf-8")
]
print("\n".join(missing))

Relevant precedent:
Existing slow PAG tests are present for SD, SD img2img/inpaint, SDXL, SDXL img2img/inpaint, and SD3 img2img.

Suggested fix:
Add at least one @slow smoke/integration class per missing pipeline using the smallest stable public checkpoint available, covering pag_scale=0 parity and pag_scale>0 execution.

Description