Modular APG

[hlky]

What does this PR do?

Adds Adaptive Projected Guidance (#9626) to Modular Diffusers with APGGuider. Notable differences to existing Guider classes are:

• we pass latents to apply_guidance so we can convert the model output to the denoised prediction (x0), APG works best on that, see Adaptive Projected Guidance #9626 (comment) from the author
• in APGGuider's set_guider we take scheduler from pipeline, this is needed for the above conversion. set_guider runs with every DenoiseStep block so scheduler will always be updated.

Code

import torch
from diffusers import StableDiffusionXLModularPipeline
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_modular import (
  StableDiffusionXLTextEncoderStep,
  StableDiffusionXLDecodeLatentsStep,
  StableDiffusionXLInputStep,
  StableDiffusionXLAutoSetTimestepsStep,
  StableDiffusionXLAutoPrepareLatentsStep,
  StableDiffusionXLAutoPrepareAdditionalConditioningStep,
  StableDiffusionXLAutoDenoiseStep,
)
from diffusers.pipelines.modular_pipeline_builder import SequentialPipelineBlocks
from diffusers.guider import APGGuider

prompt = "A 4k dslr photo of a raccoon wearing an astronaut helmet, photorealistic."


class StableDiffusionXLMainSteps(SequentialPipelineBlocks):
  block_classes = [
      StableDiffusionXLInputStep,
      StableDiffusionXLAutoSetTimestepsStep,
      StableDiffusionXLAutoPrepareLatentsStep,
      StableDiffusionXLAutoPrepareAdditionalConditioningStep,
      StableDiffusionXLAutoDenoiseStep,
  ]
  block_prefixes = [
      "input",
      "set_timesteps",
      "prepare_latents",
      "prepare_add_cond",
      "denoise",
  ]


text_encoder_workflow = StableDiffusionXLTextEncoderStep()
decoder_workflow = StableDiffusionXLDecodeLatentsStep()
sdxl_workflow = StableDiffusionXLMainSteps()
apg_guider = APGGuider()

repo = "stabilityai/stable-diffusion-xl-base-1.0"

text_encoder_workflow.add_states_from_pretrained(repo, torch_dtype=torch.float32)
decoder_workflow.add_states_from_pretrained(repo, torch_dtype=torch.float32)
sdxl_workflow.add_states_from_pretrained(repo, torch_dtype=torch.float32)

sdxl_workflow.update_states(guider=apg_guider)


print(f" text encoder workflow: {text_encoder_workflow}")
print(f" decoder workflow: {decoder_workflow}")
print(f" main sdxl workflow: {sdxl_workflow}")
print(f" prepare_latents: {sdxl_workflow.blocks['prepare_latents_step']}")

text_node = StableDiffusionXLModularPipeline()
text_node.add_blocks(text_encoder_workflow)

sdxl_node = StableDiffusionXLModularPipeline()
sdxl_node.add_blocks(sdxl_workflow)

decoder_node = StableDiffusionXLModularPipeline()
decoder_node.add_blocks(decoder_workflow)

text_state = text_node.run_blocks(prompt=prompt)

print(f" text state: {text_state}")

generator = torch.Generator().manual_seed(0)
latents_state = sdxl_node.run_blocks(
  **text_state.intermediates,
  generator=generator,
  num_inference_steps=20,
  guidance_scale=15,
  height=896,
  width=768,
  guider_kwargs={
      "adaptive_projected_guidance_momentum": -0.5,
      "adaptive_projected_guidance_rescale_factor": 15.0,
  },
)
latents = latents_state.get_intermediate("latents")

image_state = decoder_node.run_blocks(latents=latents)
image_state.get_output("images").images[0].save("apg.png")

CFG

APG

Note that the benefit of APG is more noticeable with higher guidance_scale values, we use 15 in this example.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @yiyixuxu

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

oh super cool - thank you!