293 lines
10 KiB
Python
293 lines
10 KiB
Python
# Copyright 2022 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Diffusion model utilities for prompt-to-prompt style image manipulation.
|
|
|
|
This module contains helper functions for running diffusion steps,
|
|
initializing latents, and controlling attention in diffusion models.
|
|
|
|
Based on prompt-to-prompt (https://github.com/google/prompt-to-prompt).
|
|
"""
|
|
|
|
import torch
|
|
from typing import Optional, Tuple
|
|
from einops import rearrange
|
|
|
|
|
|
def diffusion_step(
|
|
model,
|
|
controller,
|
|
latents: torch.Tensor,
|
|
context: torch.Tensor,
|
|
t: int,
|
|
guidance_scale: float,
|
|
low_resource: bool = False
|
|
) -> torch.Tensor:
|
|
"""
|
|
Perform a single diffusion step with classifier-free guidance.
|
|
|
|
Args:
|
|
model: A diffusers StableDiffusionPipeline or similar model
|
|
controller: A controller object with step_callback method
|
|
latents: Current latent tensor
|
|
context: Text embeddings context (concatenated uncond and cond)
|
|
t: Current timestep
|
|
guidance_scale: Classifier-free guidance scale
|
|
low_resource: If True, run uncond and cond separately to save memory
|
|
|
|
Returns:
|
|
Updated latent tensor after the diffusion step
|
|
"""
|
|
# Scale model input if the scheduler requires it (e.g. EulerDiscreteScheduler)
|
|
scaled_latents = model.scheduler.scale_model_input(latents, t)
|
|
if low_resource:
|
|
noise_pred_uncond = model.unet(scaled_latents, t, encoder_hidden_states=context[0])["sample"]
|
|
noise_prediction_text = model.unet(scaled_latents, t, encoder_hidden_states=context[1])["sample"]
|
|
else:
|
|
latents_input = torch.cat([scaled_latents] * 2)
|
|
noise_pred = model.unet(latents_input, t, encoder_hidden_states=context)["sample"]
|
|
noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
|
|
|
|
noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
|
|
latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"]
|
|
latents = controller.step_callback(latents)
|
|
return latents
|
|
|
|
|
|
def diffusion_step_no_cfg(
|
|
model,
|
|
controller,
|
|
latents: torch.Tensor,
|
|
context: torch.Tensor,
|
|
t: int,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Perform a single diffusion step WITHOUT classifier-free guidance.
|
|
|
|
Used for distilled models (SD Turbo, LCM) where CFG is baked into
|
|
the distillation and guidance_scale should be 0. Runs the UNet once
|
|
with only the conditional embeddings, halving memory and compute.
|
|
|
|
Args:
|
|
model: A diffusers-compatible model
|
|
controller: A controller object with step_callback method
|
|
latents: Current latent tensor [B, C, H, W] where B=2 (left, right)
|
|
context: Text embeddings [B, seq_len, dim] (cond only, no uncond)
|
|
t: Current timestep
|
|
|
|
Returns:
|
|
Updated latent tensor after the diffusion step
|
|
"""
|
|
# Scale model input if the scheduler requires it (e.g. EulerDiscreteScheduler)
|
|
latent_model_input = model.scheduler.scale_model_input(latents, t)
|
|
noise_pred = model.unet(latent_model_input, t, encoder_hidden_states=context)["sample"]
|
|
latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"]
|
|
latents = controller.step_callback(latents)
|
|
return latents
|
|
|
|
|
|
def init_latent(
|
|
latent: Optional[torch.Tensor],
|
|
model,
|
|
height: int,
|
|
width: int,
|
|
generator: Optional[torch.Generator],
|
|
batch_size: int
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Initialize or expand latents for diffusion.
|
|
|
|
Args:
|
|
latent: Optional initial latent tensor, or None to generate random
|
|
model: A diffusers model with unet attribute
|
|
height: Image height in pixels
|
|
width: Image width in pixels
|
|
generator: Optional torch generator for reproducibility
|
|
batch_size: Number of images to generate
|
|
|
|
Returns:
|
|
Tuple of (original latent, expanded latents for batch)
|
|
"""
|
|
if latent is None:
|
|
latent = torch.randn(
|
|
(1, model.unet.in_channels, height // 8, width // 8),
|
|
generator=generator,
|
|
)
|
|
latents = latent.expand(batch_size, model.unet.in_channels, height // 8, width // 8).to(model.device)
|
|
return latent, latents
|
|
|
|
|
|
def _get_unet_for_attention(model):
|
|
"""
|
|
Extract the UNet from various model wrapper types for attention registration.
|
|
|
|
Args:
|
|
model: A diffusers pipeline, ComfyUIModelWrapper, or raw model
|
|
|
|
Returns:
|
|
The UNet module
|
|
"""
|
|
# ComfyUIModelWrapper stores the wrapped UNet
|
|
if hasattr(model, 'comfy_model'):
|
|
return model.comfy_model.model.diffusion_model
|
|
|
|
# Try diffusers pipeline style
|
|
if hasattr(model, 'unet'):
|
|
return model.unet
|
|
|
|
# Try ComfyUI raw model style
|
|
if hasattr(model, 'model') and hasattr(model.model, 'diffusion_model'):
|
|
return model.model.diffusion_model
|
|
|
|
# Assume it's already a UNet
|
|
return model
|
|
|
|
|
|
def register_attention_control(model, controller):
|
|
"""
|
|
Register attention control hooks on a diffusers model.
|
|
|
|
This allows intercepting and modifying attention weights during inference.
|
|
Used for prompt-to-prompt style editing and attention visualization.
|
|
|
|
Args:
|
|
model: A diffusers StableDiffusionPipeline, ComfyUIModelWrapper, or similar model
|
|
controller: A controller object that processes attention, or None for dummy
|
|
"""
|
|
# Check if this is a ComfyUI model (has comfy_model attribute)
|
|
is_comfyui = hasattr(model, 'comfy_model')
|
|
|
|
def ca_forward_diffusers(self, place_in_unet):
|
|
"""Forward function for diffusers models."""
|
|
to_out = self.to_out
|
|
if type(to_out) is torch.nn.modules.container.ModuleList:
|
|
to_out = self.to_out[0]
|
|
else:
|
|
to_out = self.to_out
|
|
|
|
def forward(x, context=None, mask=None):
|
|
batch_size, sequence_length, dim = x.shape
|
|
h = self.heads
|
|
q = self.to_q(x)
|
|
is_cross = context is not None
|
|
context = context if is_cross else x
|
|
k = self.to_k(context)
|
|
v = self.to_v(context)
|
|
q = self.reshape_heads_to_batch_dim(q)
|
|
k = self.reshape_heads_to_batch_dim(k)
|
|
v = self.reshape_heads_to_batch_dim(v)
|
|
|
|
sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
|
|
|
|
if mask is not None:
|
|
mask = mask.reshape(batch_size, -1)
|
|
max_neg_value = -torch.finfo(sim.dtype).max
|
|
mask = mask[:, None, :].repeat(h, 1, 1)
|
|
sim.masked_fill_(~mask, max_neg_value)
|
|
|
|
# attention, what we cannot get enough of
|
|
attn = sim.softmax(dim=-1)
|
|
attn = controller(attn, is_cross, place_in_unet)
|
|
out = torch.einsum("b i j, b j d -> b i d", attn, v)
|
|
out = self.reshape_batch_dim_to_heads(out)
|
|
return to_out(out)
|
|
|
|
return forward
|
|
|
|
def ca_forward_comfyui(self, place_in_unet):
|
|
"""Forward function for ComfyUI models - accepts value and transformer_options."""
|
|
to_out = self.to_out
|
|
if type(to_out) is torch.nn.modules.container.ModuleList:
|
|
to_out = self.to_out[0]
|
|
else:
|
|
to_out = self.to_out
|
|
|
|
# ComfyUI uses dim_head instead of scale, compute scale here
|
|
scale = self.dim_head ** -0.5
|
|
|
|
def forward(x, context=None, value=None, mask=None, transformer_options=None, **kwargs):
|
|
batch_size, sequence_length, dim = x.shape
|
|
h = self.heads
|
|
q = self.to_q(x)
|
|
is_cross = context is not None
|
|
context = context if is_cross else x
|
|
k = self.to_k(context)
|
|
# Use provided value or compute from context
|
|
if value is not None:
|
|
v = self.to_v(value)
|
|
else:
|
|
v = self.to_v(context)
|
|
|
|
# Use einops rearrange instead of diffusers reshape methods
|
|
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
|
|
|
|
sim = torch.einsum("b i d, b j d -> b i j", q, k) * scale
|
|
|
|
if mask is not None:
|
|
mask = mask.reshape(batch_size, -1)
|
|
max_neg_value = -torch.finfo(sim.dtype).max
|
|
mask = mask[:, None, :].repeat(h, 1, 1)
|
|
sim.masked_fill_(~mask, max_neg_value)
|
|
|
|
# attention, what we cannot get enough of
|
|
attn = sim.softmax(dim=-1)
|
|
attn = controller(attn, is_cross, place_in_unet)
|
|
out = torch.einsum("b i j, b j d -> b i d", attn, v)
|
|
# Rearrange back to original shape
|
|
out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
|
|
return to_out(out)
|
|
|
|
return forward
|
|
|
|
# Select the appropriate forward function
|
|
ca_forward = ca_forward_comfyui if is_comfyui else ca_forward_diffusers
|
|
|
|
class DummyController:
|
|
"""Dummy controller that passes attention through unchanged."""
|
|
|
|
def __call__(self, *args):
|
|
return args[0]
|
|
|
|
def __init__(self):
|
|
self.num_att_layers = 0
|
|
|
|
if controller is None:
|
|
controller = DummyController()
|
|
|
|
def register_recr(net_, count, place_in_unet):
|
|
if net_.__class__.__name__ == 'CrossAttention':
|
|
net_.forward = ca_forward(net_, place_in_unet)
|
|
return count + 1
|
|
elif hasattr(net_, 'children'):
|
|
for net__ in net_.children():
|
|
count = register_recr(net__, count, place_in_unet)
|
|
return count
|
|
|
|
cross_att_count = 0
|
|
|
|
# Get the UNet - handle different model structures
|
|
unet = _get_unet_for_attention(model)
|
|
sub_nets = unet.named_children()
|
|
|
|
for net in sub_nets:
|
|
if "down" in net[0] or "input" in net[0]:
|
|
cross_att_count += register_recr(net[1], 0, "down")
|
|
elif "up" in net[0] or "output" in net[0]:
|
|
cross_att_count += register_recr(net[1], 0, "up")
|
|
elif "mid" in net[0]:
|
|
cross_att_count += register_recr(net[1], 0, "mid")
|
|
|
|
controller.num_att_layers = cross_att_count
|