wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
219 lines
9.6 KiB
Python
219 lines
9.6 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import torchvision.transforms as T
|
|
from einops import repeat
|
|
|
|
from invokeai.app.invocations.fields import FluxKontextConditioningField
|
|
from invokeai.app.invocations.model import VAEField
|
|
from invokeai.app.services.shared.invocation_context import InvocationContext
|
|
from invokeai.backend.flux.modules.autoencoder import AutoEncoder
|
|
from invokeai.backend.flux.sampling_utils import pack
|
|
from invokeai.backend.util.devices import TorchDevice
|
|
|
|
|
|
def generate_img_ids_with_offset(
|
|
latent_height: int,
|
|
latent_width: int,
|
|
batch_size: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
idx_offset: int = 0,
|
|
h_offset: int = 0,
|
|
w_offset: int = 0,
|
|
) -> torch.Tensor:
|
|
"""Generate tensor of image position ids with optional index and spatial offsets.
|
|
|
|
Args:
|
|
latent_height (int): Height of image in latent space (after packing, this becomes h//2).
|
|
latent_width (int): Width of image in latent space (after packing, this becomes w//2).
|
|
batch_size (int): Number of images in the batch.
|
|
device (torch.device): Device to create tensors on.
|
|
dtype (torch.dtype): Data type for the tensors.
|
|
idx_offset (int): Offset to add to the first dimension of the image ids (default: 0).
|
|
h_offset (int): Spatial offset for height/y-coordinates in latent space (default: 0).
|
|
w_offset (int): Spatial offset for width/x-coordinates in latent space (default: 0).
|
|
|
|
Returns:
|
|
torch.Tensor: Image position ids with shape [batch_size, (latent_height//2 * latent_width//2), 3].
|
|
"""
|
|
|
|
if device.type == "mps":
|
|
orig_dtype = dtype
|
|
dtype = torch.float16
|
|
|
|
# After packing, the spatial dimensions are halved due to the 2x2 patch structure
|
|
packed_height = latent_height // 2
|
|
packed_width = latent_width // 2
|
|
|
|
# Convert spatial offsets from latent space to packed space
|
|
packed_h_offset = h_offset // 2
|
|
packed_w_offset = w_offset // 2
|
|
|
|
# Create base tensor for position IDs with shape [packed_height, packed_width, 3]
|
|
# The 3 channels represent: [batch_offset, y_position, x_position]
|
|
img_ids = torch.zeros(packed_height, packed_width, 3, device=device, dtype=dtype)
|
|
|
|
# Set the batch offset for all positions
|
|
img_ids[..., 0] = idx_offset
|
|
|
|
# Create y-coordinate indices (vertical positions) with spatial offset
|
|
y_indices = torch.arange(packed_height, device=device, dtype=dtype) + packed_h_offset
|
|
# Broadcast y_indices to match the spatial dimensions [packed_height, 1]
|
|
img_ids[..., 1] = y_indices[:, None]
|
|
|
|
# Create x-coordinate indices (horizontal positions) with spatial offset
|
|
x_indices = torch.arange(packed_width, device=device, dtype=dtype) + packed_w_offset
|
|
# Broadcast x_indices to match the spatial dimensions [1, packed_width]
|
|
img_ids[..., 2] = x_indices[None, :]
|
|
|
|
# Expand to include batch dimension: [batch_size, (packed_height * packed_width), 3]
|
|
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=batch_size)
|
|
|
|
if device.type == "mps":
|
|
img_ids = img_ids.to(orig_dtype)
|
|
|
|
return img_ids
|
|
|
|
|
|
class KontextExtension:
|
|
"""Applies FLUX Kontext (reference image) conditioning."""
|
|
|
|
def __init__(
|
|
self,
|
|
kontext_conditioning: list[FluxKontextConditioningField],
|
|
context: InvocationContext,
|
|
vae_field: VAEField,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
):
|
|
"""
|
|
Initializes the KontextExtension, pre-processing the reference images
|
|
into latents and positional IDs.
|
|
"""
|
|
self._context = context
|
|
self._device = device
|
|
self._dtype = dtype
|
|
self._vae_field = vae_field
|
|
self.kontext_conditioning = kontext_conditioning
|
|
|
|
# Pre-process and cache the kontext latents and ids upon initialization.
|
|
self.kontext_latents, self.kontext_ids = self._prepare_kontext()
|
|
|
|
def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Encodes the reference images and prepares their concatenated latents and IDs with spatial tiling."""
|
|
all_latents = []
|
|
all_ids = []
|
|
|
|
# Track cumulative dimensions for spatial tiling
|
|
# These track the running extent of the virtual canvas in latent space
|
|
canvas_h = 0 # Running canvas height
|
|
canvas_w = 0 # Running canvas width
|
|
|
|
vae_info = self._context.models.load(self._vae_field.vae)
|
|
|
|
for idx, kontext_field in enumerate(self.kontext_conditioning):
|
|
image = self._context.images.get_pil(kontext_field.image.image_name)
|
|
|
|
# Convert to RGB
|
|
image = image.convert("RGB")
|
|
|
|
# Convert to tensor using torchvision transforms for consistency
|
|
transformation = T.Compose(
|
|
[
|
|
T.ToTensor(), # Converts PIL image to tensor and scales to [0, 1]
|
|
]
|
|
)
|
|
image_tensor = transformation(image)
|
|
# Convert from [0, 1] to [-1, 1] range expected by VAE
|
|
image_tensor = image_tensor * 2.0 - 1.0
|
|
image_tensor = image_tensor.unsqueeze(0) # Add batch dimension
|
|
image_tensor = image_tensor.to(self._device)
|
|
|
|
# Continue with VAE encoding
|
|
# Don't sample from the distribution for reference images - use the mean (matching ComfyUI)
|
|
# Estimate working memory for encode operation (50% of decode memory requirements)
|
|
img_h = image_tensor.shape[-2]
|
|
img_w = image_tensor.shape[-1]
|
|
element_size = next(vae_info.model.parameters()).element_size()
|
|
scaling_constant = 1100 # 50% of decode scaling constant (2200)
|
|
estimated_working_memory = int(img_h * img_w * element_size * scaling_constant)
|
|
|
|
with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
|
|
assert isinstance(vae, AutoEncoder)
|
|
vae_dtype = next(iter(vae.parameters())).dtype
|
|
image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
|
|
# Use sample=False to get the distribution mean without noise
|
|
kontext_latents_unpacked = vae.encode(image_tensor, sample=False)
|
|
TorchDevice.empty_cache()
|
|
|
|
# Extract tensor dimensions
|
|
batch_size, _, latent_height, latent_width = kontext_latents_unpacked.shape
|
|
|
|
# Pad latents to be compatible with patch_size=2
|
|
# This ensures dimensions are even for the pack() function
|
|
pad_h = (2 - latent_height % 2) % 2
|
|
pad_w = (2 - latent_width % 2) % 2
|
|
if pad_h > 0 or pad_w > 0:
|
|
kontext_latents_unpacked = F.pad(kontext_latents_unpacked, (0, pad_w, 0, pad_h), mode="circular")
|
|
# Update dimensions after padding
|
|
_, _, latent_height, latent_width = kontext_latents_unpacked.shape
|
|
|
|
# Pack the latents
|
|
kontext_latents_packed = pack(kontext_latents_unpacked).to(self._device, self._dtype)
|
|
|
|
# Determine spatial offsets for this reference image
|
|
h_offset = 0
|
|
w_offset = 0
|
|
|
|
if idx > 0: # First image starts at (0, 0)
|
|
# Calculate potential canvas dimensions for each tiling option
|
|
# Option 1: Tile vertically (below existing content)
|
|
potential_h_vertical = canvas_h + latent_height
|
|
|
|
# Option 2: Tile horizontally (to the right of existing content)
|
|
potential_w_horizontal = canvas_w + latent_width
|
|
|
|
# Choose arrangement that minimizes the maximum dimension
|
|
# This keeps the canvas closer to square, optimizing attention computation
|
|
if potential_h_vertical > potential_w_horizontal:
|
|
# Tile horizontally (to the right of existing images)
|
|
w_offset = canvas_w
|
|
canvas_w = canvas_w + latent_width
|
|
canvas_h = max(canvas_h, latent_height)
|
|
else:
|
|
# Tile vertically (below existing images)
|
|
h_offset = canvas_h
|
|
canvas_h = canvas_h + latent_height
|
|
canvas_w = max(canvas_w, latent_width)
|
|
else:
|
|
# First image - just set canvas dimensions
|
|
canvas_h = latent_height
|
|
canvas_w = latent_width
|
|
|
|
# Generate IDs with both index offset and spatial offsets
|
|
kontext_ids = generate_img_ids_with_offset(
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
batch_size=batch_size,
|
|
device=self._device,
|
|
dtype=self._dtype,
|
|
idx_offset=1, # All reference images use index=1 (matching ComfyUI implementation)
|
|
h_offset=h_offset,
|
|
w_offset=w_offset,
|
|
)
|
|
|
|
all_latents.append(kontext_latents_packed)
|
|
all_ids.append(kontext_ids)
|
|
|
|
# Concatenate all latents and IDs along the sequence dimension
|
|
concatenated_latents = torch.cat(all_latents, dim=1) # Concatenate along sequence dimension
|
|
concatenated_ids = torch.cat(all_ids, dim=1) # Concatenate along sequence dimension
|
|
|
|
return concatenated_latents, concatenated_ids
|
|
|
|
def ensure_batch_size(self, target_batch_size: int) -> None:
|
|
"""Ensures the kontext latents and IDs match the target batch size by repeating if necessary."""
|
|
if self.kontext_latents.shape[0] != target_batch_size:
|
|
self.kontext_latents = self.kontext_latents.repeat(target_batch_size, 1, 1)
|
|
self.kontext_ids = self.kontext_ids.repeat(target_batch_size, 1, 1)
|