Implementation:Datajuicer Data juicer Prompt2PromptPipeline

Knowledge Sources	Datajuicer_Data_juicer
Domains	Generative AI, Image Editing, Stable Diffusion, Attention Control
Last Updated	2026-02-14 16:00 GMT

Overview

Implements the Prompt-to-Prompt image editing pipeline for Stable Diffusion XL, enabling fine-grained control over generated images by manipulating cross-attention maps during the diffusion process.

Description

This module is adapted from the Prompt-to-Prompt with SDXL project and provides sophisticated text-driven image editing capabilities within the Data-Juicer pipeline. It powers the SDXL Prompt-to-Prompt mapper operator.

Core Pipeline:

• Prompt2PromptPipeline -- Extends `StableDiffusionXLPipeline` with attention controller injection. During the denoising loop, it intercepts cross-attention computations via `P2PCrossAttnProcessor` to enable three edit types: replace, refine, and reweight.

Attention Control Hierarchy:

• AttentionControl (ABC) -- Base class managing attention layer counting and step tracking.
• AttentionStore -- Captures and accumulates attention maps across diffusion steps for analysis.
• AttentionControlEdit (ABC) -- Adds self-attention and cross-attention replacement logic with configurable replacement schedules.
• AttentionReplace -- Swaps attention maps for word replacement using token-level mapper matrices.
• AttentionRefine -- Adjusts attention for prompt refinement using global sequence alignment (Needleman-Wunsch).
• AttentionReweight -- Scales attention weights for emphasis control using per-token equalizer values.

Spatial Control:

• LocalBlend -- Applies spatial masking to constrain edits to specific image regions, using one-hot alpha layers based on specified blend words.

Cross-Attention Processor:

• P2PCrossAttnProcessor -- Custom attention processor that intercepts attention probability computation and feeds it to the controller.

Utility Functions:

• create_controller() -- Factory function creating the appropriate attention controller based on edit type and parameters.
• get_word_inds() -- Maps words to token indices accounting for subword tokenization.
• get_replacement_mapper() -- Creates token-to-token mapping matrices for replacement edits.
• get_refinement_mapper() -- Creates aligned token mappings using Needleman-Wunsch global alignment.
• get_equalizer() -- Creates per-token weight tensors for reweight edits.

Usage

Used internally by the SDXL Prompt-to-Prompt mapper operator. Requires a Stable Diffusion XL model and appropriate prompts with edit configuration.

Code Reference

Source Location

• Repository: Datajuicer_Data_juicer
• File: data_juicer/ops/common/prompt2prompt_pipeline.py
• Lines: 1-1215

Signature

class Prompt2PromptPipeline(StableDiffusionXLPipeline):
    @torch.no_grad()
    def __call__(self, prompt, prompt_2=None, height=None, width=None,
                 num_inference_steps=50, guidance_scale=7.5,
                 cross_attention_kwargs=None, attn_res=None, ...): ...
    def register_attention_control(self, controller): ...

class AttentionReplace(AttentionControlEdit):
    def replace_cross_attention(self, attn_base, att_replace): ...

class AttentionRefine(AttentionControlEdit):
    def replace_cross_attention(self, attn_base, att_replace): ...

class AttentionReweight(AttentionControlEdit):
    def replace_cross_attention(self, attn_base, att_replace): ...

class LocalBlend:
    def __call__(self, x_t, attention_store): ...

def create_controller(prompts, cross_attention_kwargs, num_inference_steps,
                      tokenizer, device, attn_res) -> AttentionControl: ...

Import

from data_juicer.ops.common.prompt2prompt_pipeline import (
    Prompt2PromptPipeline,
    AttentionReplace,
    AttentionRefine,
    AttentionReweight,
    LocalBlend,
    create_controller,
)

I/O Contract

Inputs

Name	Type	Required	Description
prompt	str or List[str]	Yes	Source and target prompts for image generation/editing
cross_attention_kwargs	Dict	Yes	Edit configuration: edit_type ("replace"/"refine"/"reweight"), n_cross_replace, n_self_replace, local_blend_words, equalizer_words, equalizer_strengths
height	int	No	Image height in pixels (default: model default)
width	int	No	Image width in pixels (default: model default)
num_inference_steps	int	No	Number of denoising steps (default: 50)
guidance_scale	float	No	Classifier-free guidance scale (default: 7.5)

Outputs

Name	Type	Description
images	StableDiffusionXLPipelineOutput	Generated/edited images as PIL Images or tensors

Usage Examples

from data_juicer.ops.common.prompt2prompt_pipeline import Prompt2PromptPipeline

# Load pipeline (requires SDXL model)
pipe = Prompt2PromptPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")

# Word replacement edit
result = pipe(
    prompt=["a photo of a cat sitting on a bench", "a photo of a dog sitting on a bench"],
    cross_attention_kwargs={
        "edit_type": "replace",
        "n_cross_replace": 0.4,
        "n_self_replace": 0.4,
        "local_blend_words": [["cat", "dog"]],
    },
    num_inference_steps=50,
)

# Attention reweighting
result = pipe(
    prompt=["a photo of a big red car", "a photo of a big red car"],
    cross_attention_kwargs={
        "edit_type": "reweight",
        "equalizer_words": ["big"],
        "equalizer_strengths": [2.0],
    },
)

Related Pages

• data_juicer/ops/mapper/ (SDXL Prompt-to-Prompt mapper)
• Datajuicer_Data_juicer_Operator_Base_Classes