seamless_dit.py

"""Seamless (tileable) generation for transformer / DiT diffusion models.

The classic seamless-tiling trick (switching every Conv2d to circular padding)
only works on convolutional U-Nets (SD1.5 / SDXL). Transformer models such as
Flux and Wan/Anima patchify with a Linear or a non-overlapping Conv, so there is
nothing for that trick to grab.

This pack provides a model-agnostic alternative:

* ``SeamlessTileModelDiT`` rolls the latent by a per-step offset during sampling
  (via a UNet function wrapper) so no fixed seam can form. Works on Flux, Wan,
  and SD/SDXL alike, and never mutates ComfyUI's cached model.
* ``MakeCircularVAEDiT`` applies circular padding to the VAE decoder, handling
  both Conv2d (SD/Flux AutoencoderKL) and Conv3d / CausalConv3d (Wan VAE) for
  pixel-perfect edges.
"""

import copy

import torch
from torch import Tensor
from torch.nn import functional as F
from torch.nn.modules.utils import _pair, _triple


TILING_MODES = ["enable", "x_only", "y_only", "disable"]


def _axes_for(tiling):
    """Return (tileX, tileY): X = wrap horizontally (width), Y = vertically."""
    return (
        tiling in ("enable", "x_only"),  # tileX -> width
        tiling in ("enable", "y_only"),  # tileY -> height
    )


# ---------------------------------------------------------------------------
# Node 1: model-side seamless via latent rolling
# ---------------------------------------------------------------------------

def _step_offsets(seed, timestep, height, width, tile_x, tile_y, scale=1.0):
    """Deterministic per-step roll offsets, scaled by ``scale`` (0..1).

    Keyed on the timestep value so that, within one sampling step, the cond and
    uncond passes get the SAME offset (otherwise CFG would combine misaligned
    predictions), while different steps get different offsets.

    ``scale`` tapers the offset magnitude with the noise level: full-range rolls
    at high sigma establish tileable structure, while the low-sigma detail steps
    barely roll so their wrap-artifact collapses onto the true edge (handled by
    the circular VAE) instead of leaving a faint seam inside the image.
    """
    try:
        t_val = float(timestep.flatten()[0].item())
    except Exception:
        t_val = 0.0
    key = (int(seed) & 0x7FFFFFFF) ^ (int(t_val * 100000.0) & 0x7FFFFFFF)
    gen = torch.Generator().manual_seed(int(key))
    dy = int(round(scale * torch.randint(0, max(1, height), (1,), generator=gen).item())) if tile_y else 0
    dx = int(round(scale * torch.randint(0, max(1, width), (1,), generator=gen).item())) if tile_x else 0
    return dy, dx


def _roll_conds(c, shifts, dims):
    """Roll spatial conditioning (c_concat) to stay aligned with the latent.

    Text conditioning (c_crossattn) is not spatial and is left alone; control
    hints / transformer_options are out of scope (see module docstring)."""
    cc = c.get("c_concat", None)
    if cc is None or not torch.is_tensor(cc):
        return c
    new_c = dict(c)
    new_c["c_concat"] = torch.roll(cc, shifts=shifts, dims=dims)
    return new_c


def _make_tiling_wrapper(seed, tiling):
    tile_x, tile_y = _axes_for(tiling)
    # Per-run state: track the largest sigma seen so we can normalise the current
    # noise level to [0, 1] and taper the roll magnitude toward the end.
    state = {"max_sigma": 0.0}

    def wrapper(apply_model, params):
        inp = params["input"]
        timestep = params["timestep"]
        c = params["c"]

        if not (tile_x or tile_y):
            return apply_model(inp, timestep, **c)

        try:
            sigma = float(timestep.flatten().max().item())
        except Exception:
            sigma = 0.0
        if sigma > state["max_sigma"]:
            state["max_sigma"] = sigma
        scale = (sigma / state["max_sigma"]) if state["max_sigma"] > 0 else 1.0

        height, width = inp.shape[-2], inp.shape[-1]
        dy, dx = _step_offsets(seed, timestep, height, width, tile_x, tile_y, scale)
        if dy == 0 and dx == 0:
            return apply_model(inp, timestep, **c)

        shifts = (dy, dx)
        dims = (-2, -1)  # last two dims are H, W for both 4D and 5D latents
        rolled = torch.roll(inp, shifts=shifts, dims=dims)
        c_rolled = _roll_conds(c, shifts, dims)
        out = apply_model(rolled, timestep, **c_rolled)
        # Un-roll so the sampler's canonical latent frame is preserved.
        return torch.roll(out, shifts=(-dy, -dx), dims=dims)

    return wrapper


class SeamlessTileModelDiT:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "tiling": (TILING_MODES,),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "run"
    CATEGORY = "conditioning"

    def run(self, model, tiling, seed):
        # clone() gives us a private model_options dict; the shared/cached model
        # weights are never touched, so this cannot poison other generations.
        m = model.clone()
        if tiling == "disable":
            return (m,)
        m.set_model_unet_function_wrapper(_make_tiling_wrapper(seed, tiling))
        return (m,)


# ---------------------------------------------------------------------------
# Node 2: circular VAE decode (Conv2d + Conv3d / CausalConv3d)
# ---------------------------------------------------------------------------

def _replacement_conv2d_forward(self, input: Tensor, weight: Tensor, bias):
    working = F.pad(input, self._tile_padX, mode=self._tile_modeX)
    working = F.pad(working, self._tile_padY, mode=self._tile_modeY)
    return F.conv2d(working, weight, bias, self.stride, _pair(0), self.dilation, self.groups)


def _patch_conv2d(layer, tile_x, tile_y):
    rp = layer._reversed_padding_repeated_twice
    layer._tile_modeX = "circular" if tile_x else "constant"
    layer._tile_modeY = "circular" if tile_y else "constant"
    layer._tile_padX = (rp[0], rp[1], 0, 0)
    layer._tile_padY = (0, 0, rp[2], rp[3])
    layer._conv_forward = _replacement_conv2d_forward.__get__(layer, layer.__class__)


def _replacement_conv3d_forward(self, input, weight, bias, autopad=None, *args, **kwargs):
    # Circular/constant pad only the spatial dims (H, W); the temporal dim is
    # handled by the (already-run) causal padding logic and by the delegated
    # conv via the temporal padding we left on self.padding.
    if self._tile_pw > 0:
        input = F.pad(input, (self._tile_pw, self._tile_pw, 0, 0, 0, 0), mode=self._tile_modeX)
    if self._tile_ph > 0:
        input = F.pad(input, (0, 0, self._tile_ph, self._tile_ph, 0, 0), mode=self._tile_modeY)
    # Only forward the `autopad` kwarg when it's set: plain torch Conv3d doesn't
    # accept it, while ComfyUI's ops.Conv3d/CausalConv3d do (and use it for the
    # single-frame causal fast path).
    if autopad is not None:
        return self._tile_orig_conv_forward(input, weight, bias, autopad=autopad, *args, **kwargs)
    return self._tile_orig_conv_forward(input, weight, bias, *args, **kwargs)


def _patch_conv3d(layer, tile_x, tile_y):
    if not getattr(layer, "_tile_wrapped", False):
        pd, ph, pw = _triple(layer.padding)
        if ph == 0 and pw == 0:
            return  # nothing spatial to wrap (e.g. 1x1x1 projection convs)
        layer._tile_ph, layer._tile_pw = ph, pw
        layer._tile_orig_conv_forward = layer._conv_forward
        # Disable the layer's own spatial padding; we apply it manually above so
        # circular wrapping works and it isn't double-padded. Temporal padding
        # (pd) is preserved for the delegated conv.
        layer.padding = (pd, 0, 0)
        layer._conv_forward = _replacement_conv3d_forward.__get__(layer, layer.__class__)
        layer._tile_wrapped = True
    layer._tile_modeX = "circular" if tile_x else "constant"
    layer._tile_modeY = "circular" if tile_y else "constant"


def make_circular(module, tile_x, tile_y):
    """Apply circular (or constant) spatial padding to every conv in a module.

    Handles Conv2d (SD/Flux VAE) and Conv3d/CausalConv3d (Wan VAE)."""
    for layer in module.modules():
        if isinstance(layer, torch.nn.Conv2d):
            _patch_conv2d(layer, tile_x, tile_y)
        elif isinstance(layer, torch.nn.Conv3d):
            _patch_conv3d(layer, tile_x, tile_y)
    return module


class MakeCircularVAEDiT:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vae": ("VAE",),
                "tiling": (TILING_MODES,),
                "copy_vae": (["Make a copy", "Modify in place"],),
            },
        }

    RETURN_TYPES = ("VAE",)
    FUNCTION = "run"
    CATEGORY = "latent"

    def run(self, vae, tiling, copy_vae):
        if copy_vae == "Modify in place":
            vae_copy = vae
        else:
            # deepcopy isolates from the cached VAE; the internal patcher and
            # first_stage_model stay the same (shared) object within the copy,
            # so mutating first_stage_model is what actually gets decoded.
            vae_copy = copy.deepcopy(vae)

        tile_x, tile_y = _axes_for(tiling)
        make_circular(vae_copy.first_stage_model, tile_x, tile_y)
        return (vae_copy,)


NODE_CLASS_MAPPINGS = {
    "SeamlessTileModelDiT": SeamlessTileModelDiT,
    "MakeCircularVAEDiT": MakeCircularVAEDiT,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SeamlessTileModelDiT": "Seamless Tile Model (DiT)",
    "MakeCircularVAEDiT": "Make Circular VAE (DiT)",
}