Prompt scheduling in Diffusers like A1111

[linhbeige]

Hi everyone, I have a question that how to implement the prompt scheduling feature in A1111 by diffusers library.

Example prompt: Official portrait of a smiling world war ii general, [male:female:0.99], cheerful, happy, detailed face, 20th century, highly detailed, cinematic lighting, digital art painting by Greg Rutkowski.

[asomoza]

HI, there isn't a built-in option to do this, so you'll have to implement yourself. The easiest way to achieve this is using the callbacks.

Basically you have to encode the different prompts beforehand and change them on the steps you want.

[hlky]

Example for StableDiffusionPipeline, two versions, one that encodes the prompt inside the callback and the other for encoding the prompt beforehand.

from diffusers import StableDiffusionPipeline
from diffusers.callbacks import PipelineCallback, MultiPipelineCallbacks
from diffusers.configuration_utils import register_to_config
import torch
from typing import Any, Dict, Optional, Tuple, Union


class SDPromptScheduleCallback(PipelineCallback):
    @register_to_config
    def __init__(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        num_images_per_prompt: int = 1,
        cutoff_step_ratio=1.0,
        cutoff_step_index=None,
    ):
        super().__init__(
            cutoff_step_ratio=cutoff_step_ratio, cutoff_step_index=cutoff_step_index
        )

    tensor_inputs = ["prompt_embeds"]

    def callback_fn(
        self, pipeline, step_index, timestep, callback_kwargs
    ) -> Dict[str, Any]:
        cutoff_step_ratio = self.config.cutoff_step_ratio
        cutoff_step_index = self.config.cutoff_step_index

        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
        cutoff_step = (
            cutoff_step_index
            if cutoff_step_index is not None
            else int(pipeline.num_timesteps * cutoff_step_ratio)
        )

        if step_index == cutoff_step:
            prompt_embeds, negative_prompt_embeds = pipeline.encode_prompt(
                prompt=self.config.prompt,
                negative_prompt=self.config.negative_prompt,
                device=pipeline._execution_device,
                num_images_per_prompt=self.config.num_images_per_prompt,
                do_classifier_free_guidance=pipeline.do_classifier_free_guidance,
            )
            if pipeline.do_classifier_free_guidance:
                prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
            callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
        return callback_kwargs


class SDPromptSchedulingCallback(PipelineCallback):
    @register_to_config
    def __init__(
        self,
        encoded_prompt: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        cutoff_step_ratio=1.0,
        cutoff_step_index=None,
    ):
        super().__init__(
            cutoff_step_ratio=cutoff_step_ratio, cutoff_step_index=cutoff_step_index
        )

    tensor_inputs = ["prompt_embeds"]

    def callback_fn(
        self, pipeline, step_index, timestep, callback_kwargs
    ) -> Dict[str, Any]:
        cutoff_step_ratio = self.config.cutoff_step_ratio
        cutoff_step_index = self.config.cutoff_step_index
        if isinstance(self.config.prompt_embeds, tuple):
            prompt_embeds, negative_prompt_embeds = self.config.encoded_prompt
        else:
            prompt_embeds = self.config.encoded_prompt

        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
        cutoff_step = (
            cutoff_step_index
            if cutoff_step_index is not None
            else int(pipeline.num_timesteps * cutoff_step_ratio)
        )

        if step_index == cutoff_step:
            if pipeline.do_classifier_free_guidance:
                prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
            callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
        return callback_kwargs


pipeline: StableDiffusionPipeline = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
).to("cuda")
pipeline.safety_checker = None
pipeline.requires_safety_checker = False

callback = MultiPipelineCallbacks(
    [
        SDPromptSchedulingCallback(
            encoded_prompt=pipeline.encode_prompt(
                prompt="Official portrait of a smiling world war ii general, female, cheerful, happy, detailed face, 20th century, highly detailed, cinematic lighting, digital art painting by Greg Rutkowski",
                negative_prompt="Deformed, ugly, bad anatomy",
                device=pipeline._execution_device,
                num_images_per_prompt=1,
                # pipeline.do_classifier_free_guidance can't be accessed until after pipeline is ran
                do_classifier_free_guidance=True,
            ),
            cutoff_step_ratio=0.25,
        )
    ]
)

callback = MultiPipelineCallbacks(
    [
        SDPromptScheduleCallback(
            prompt="Official portrait of a smiling world war ii general, female, cheerful, happy, detailed face, 20th century, highly detailed, cinematic lighting, digital art painting by Greg Rutkowski",
            negative_prompt="Deformed, ugly, bad anatomy",
            num_images_per_prompt=1,
            cutoff_step_ratio=0.25,
        )
    ]
)

image = pipeline(
    prompt="Official portrait of a smiling world war ii general, male, cheerful, happy, detailed face, 20th century, highly detailed, cinematic lighting, digital art painting by Greg Rutkowski",
    negative_prompt="Deformed, ugly, bad anatomy",
    callback_on_step_end=callback,
    callback_on_step_end_tensor_inputs=["prompt_embeds"],
).images[0]

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[asomoza]

@hlky would you be interested in adding this in the community scripts? I think this is really good as an example for that section.

I was thinking on maybe adding it to the official callbacks but I don't think this gets used that much (I could be wrong).

[hlky]

@asomoza Sure, I've created a PR adding this to community scripts. I'm not sure how much this feature gets used either.