avtc
01a324bb11
Extends tiling beyond UNet (Conv2d) models to support DiT architectures using toroidal attention patches and latent content wrapping for seamless infinite tiling. Adds rectangular tiling mode, WanVAE 5D tensor handling, and an AdvancedTilingRay node for distributed Raylight workers.
108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
"""
|
|
DiT model tiling via toroidal attention and latent wrapping.
|
|
|
|
For Hexagon mode:
|
|
- Toroidal attention: injects wrapped-neighbor K/V entries into every attention
|
|
layer, making boundary patches structurally "see" opposite-edge content as
|
|
spatially adjacent.
|
|
- Latent wrapping: copies source content to waste positions in the input latent
|
|
on each denoising step (for KSampler preview).
|
|
|
|
For Rectangular mode:
|
|
- Toroidal attention: injects wrapped K/V entries from opposite edges (right↔left,
|
|
top↔bottom).
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import Conv2d
|
|
|
|
from .modes import Settings
|
|
|
|
|
|
def _has_conv2d(model: nn.Module) -> bool:
|
|
"""Check if the model has any Conv2d layers (UNet vs DiT)."""
|
|
return any(isinstance(m, Conv2d) for m in model.modules())
|
|
|
|
|
|
def _create_content_wrapper(settings: Settings):
|
|
"""
|
|
Create a model function wrapper that copies source content to waste
|
|
positions in the latent on each denoising step. This makes the KSampler
|
|
preview show wrapped content instead of noise.
|
|
|
|
Only used for Hexagon mode.
|
|
|
|
:param settings: Tiling settings
|
|
"""
|
|
from .advanced_tiling import calculate_mapping
|
|
|
|
_mapping_cache = {}
|
|
|
|
def wrapper(apply_model, args):
|
|
x = args["input"]
|
|
is_5d = x.ndim == 5
|
|
|
|
if is_5d:
|
|
_, _, _, H, W = x.shape
|
|
else:
|
|
_, _, H, W = x.shape
|
|
|
|
cache_key = (W, H, hash(settings))
|
|
if cache_key not in _mapping_cache:
|
|
_mapping_cache[cache_key] = calculate_mapping(
|
|
(W, H), (W, H), settings
|
|
)
|
|
mapping = _mapping_cache[cache_key]
|
|
|
|
# Content replacement in latent
|
|
if is_5d:
|
|
x[:, :, :, mapping[1], mapping[0]] = x[:, :, :, mapping[3], mapping[2]]
|
|
else:
|
|
x[:, :, mapping[1], mapping[0]] = x[:, :, mapping[3], mapping[2]]
|
|
|
|
return apply_model(args["input"], args["timestep"], **args["c"])
|
|
|
|
return wrapper
|
|
|
|
|
|
def patch_dit_model(model_patcher, settings: Settings):
|
|
"""
|
|
Apply tiling to a DiT model.
|
|
|
|
For Hexagon mode: applies toroidal attention + latent content wrapping.
|
|
For Rectangular mode: applies toroidal attention only.
|
|
|
|
:param model_patcher: ComfyUI ModelPatcher instance
|
|
:param settings: Tiling settings
|
|
"""
|
|
diff_model = model_patcher.model.diffusion_model
|
|
|
|
if settings.mode == "Hexagon":
|
|
from .toroidal_attention import HexToroidalAttentionPatch
|
|
|
|
if not hasattr(diff_model, 'pe_embedder'):
|
|
raise ValueError(
|
|
"Model does not have pe_embedder. "
|
|
"Toroidal attention requires a model with RoPE position embeddings."
|
|
)
|
|
|
|
patch = HexToroidalAttentionPatch(settings, diff_model.pe_embedder)
|
|
model_patcher.set_model_attn1_patch(patch)
|
|
|
|
# Latent content wrapping for KSampler preview
|
|
wrapper = _create_content_wrapper(settings)
|
|
model_patcher.set_model_unet_function_wrapper(wrapper)
|
|
|
|
elif settings.mode == "Rectangular":
|
|
from .toroidal_attention import RectToroidalAttentionPatch
|
|
|
|
if not hasattr(diff_model, 'pe_embedder'):
|
|
raise ValueError(
|
|
"Model does not have pe_embedder. "
|
|
"Toroidal attention requires a model with RoPE position embeddings."
|
|
)
|
|
|
|
patch = RectToroidalAttentionPatch(diff_model.pe_embedder)
|
|
model_patcher.set_model_attn1_patch(patch)
|