| """Pieces shared by the Krea 2 and Qwen-Image 2.1 UltraTex backbones. |
| |
| Both backbones are single-stream DiTs that consume one flat token sequence with |
| per-token 3-axis RoPE ids. UltraTex's Background Token Dropping (BTD) is a pure |
| sequence operation (gather the foreground tokens and their ids), so the only |
| backbone-specific parts are the block, the id convention and the modulation. |
| This module holds the attention dispatch (dense SDPA or the vendored SLA |
| block-sparse kernel), RoPE helpers, the flow-matching schedules and the |
| ComfyUI-backed VAE / latent normalisation used by both. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| import os |
| import sys |
| from dataclasses import dataclass |
|
|
| import torch |
| from einops import rearrange |
| from torch import Tensor, nn |
| from torch.nn import functional as F |
|
|
| |
| |
| |
|
|
| _SPARSE_ATTENTION = None |
| _SPARSE_ATTENTION_CONFIG: dict = {"enabled": False} |
|
|
|
|
| def configure_sparse_attention( |
| enabled: bool, |
| head_dim: int = 128, |
| topk: float = 0.20, |
| blkq: int = 128, |
| blkk: int = 64, |
| ) -> dict: |
| """Select the attention backend used by every DiT block. |
| |
| Mirrors ``ultratex/backbones/flux2/model.py``. The SLA Triton kernel needs |
| equal query/key lengths, which holds for the pruned [text, target, refs] |
| sequence that the trainers build. |
| """ |
| global _SPARSE_ATTENTION, _SPARSE_ATTENTION_CONFIG |
| if not enabled: |
| _SPARSE_ATTENTION = None |
| _SPARSE_ATTENTION_CONFIG = {"enabled": False} |
| return dict(_SPARSE_ATTENTION_CONFIG) |
| if not 0.0 < topk <= 1.0: |
| raise ValueError(f"sparse attention topk must be in (0, 1], got {topk}") |
| if head_dim != 128: |
| raise ValueError(f"The vendored SLA kernel expects head_dim=128, got {head_dim}") |
| from ultratex.sparse_attention import SparseLinearAttention |
|
|
| _SPARSE_ATTENTION = SparseLinearAttention( |
| head_dim=head_dim, topk=topk, feature_map="softmax", BLKQ=blkq, BLKK=blkk |
| ) |
| _SPARSE_ATTENTION_CONFIG = { |
| "enabled": True, |
| "backend": "SparseLinearAttention", |
| "head_dim": head_dim, |
| "topk": topk, |
| "BLKQ": blkq, |
| "BLKK": blkk, |
| } |
| return dict(_SPARSE_ATTENTION_CONFIG) |
|
|
|
|
| def sparse_attention_config() -> dict: |
| return dict(_SPARSE_ATTENTION_CONFIG) |
|
|
|
|
| def sparse_attention_enabled() -> bool: |
| return _SPARSE_ATTENTION is not None |
|
|
|
|
| def dense_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor | None = None) -> Tensor: |
| """q, k, v: [B, H, N, D] -> [B, N, H*D].""" |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) |
| return rearrange(out, "b h n d -> b n (h d)") |
|
|
|
|
| def full_attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor: |
| """Bidirectional attention over the whole sequence, sparse when configured.""" |
| if _SPARSE_ATTENTION is not None and q.shape[2] == k.shape[2]: |
| return rearrange(_SPARSE_ATTENTION(q, k, v), "b h n d -> b n (h d)") |
| return dense_attention(q, k, v) |
|
|
|
|
| def block_causal_attention(q: Tensor, k: Tensor, v: Tensor, segments: list[tuple[int, int, bool]]) -> Tensor: |
| """Qwen-Image 2.1 block-causal attention. |
| |
| ``segments`` lists ``(start, end, causal)`` spans covering the sequence in |
| order. A span attends to every key before its own end; a causal span (text) |
| additionally applies a causal triangle inside itself. The last span is the |
| target image, which therefore sees everything. |
| |
| With the sparse backend, the target rows come from one SLA pass over the |
| full sequence (the only large block) and the prefix rows are recomputed |
| exactly with dense SDPA over their restricted keys. |
| """ |
| n = q.shape[2] |
| if not segments or segments[-1][1] != n: |
| raise ValueError(f"segments {segments} do not cover the sequence of length {n}") |
| if _SPARSE_ATTENTION is not None: |
| out = rearrange(_SPARSE_ATTENTION(q, k, v), "b h n d -> b n (h d)") |
| prefix_end = segments[-1][0] |
| if prefix_end == 0: |
| return out |
| parts = [ |
| _segment_attention(q, k, v, start, end, causal) for start, end, causal in segments[:-1] |
| ] |
| return torch.cat(parts + [out[:, prefix_end:]], dim=1) |
| parts = [_segment_attention(q, k, v, start, end, causal) for start, end, causal in segments] |
| return torch.cat(parts, dim=1) if len(parts) > 1 else parts[0] |
|
|
|
|
| def _segment_attention(q, k, v, start, end, causal): |
| mask = None |
| if causal: |
| mask = torch.ones((end - start, end), dtype=torch.bool, device=q.device).tril(start) |
| return dense_attention(q[:, :, start:end], k[:, :, :end], v[:, :, :end], mask) |
|
|
|
|
| def prefix_cached_attention(q: Tensor, k: Tensor, v: Tensor, prefix_k: Tensor, prefix_v: Tensor) -> Tensor: |
| """Target-only queries against [cached prefix, target] keys (dense).""" |
| return dense_attention(q, torch.cat([prefix_k, k], dim=2), torch.cat([prefix_v, v], dim=2)) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def rope(pos: Tensor, dim: int, theta: float) -> Tensor: |
| assert dim % 2 == 0 |
| scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim |
| omega = 1.0 / (theta**scale) |
| out = torch.einsum("...n,d->...nd", pos.to(torch.float64), omega) |
| out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) |
| out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2) |
| return out.float() |
|
|
|
|
| class EmbedND(nn.Module): |
| def __init__(self, dim: int, theta: float, axes_dim: list[int]): |
| super().__init__() |
| self.dim = dim |
| self.theta = theta |
| self.axes_dim = axes_dim |
|
|
| def forward(self, ids: Tensor) -> Tensor: |
| emb = torch.cat( |
| [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(len(self.axes_dim))], dim=-3 |
| ) |
| return emb.unsqueeze(1) |
|
|
|
|
| def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: |
| xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) |
| xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) |
| xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] |
| xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] |
| return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) |
|
|
|
|
| def timestep_embedding(t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0) -> Tensor: |
| t = time_factor * t |
| half = dim // 2 |
| freqs = torch.exp(-math.log(max_period) * torch.arange(0, half, device=t.device, dtype=torch.float32) / half) |
| args = t[:, None].float() * freqs[None] |
| embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| if dim % 2: |
| embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| if torch.is_floating_point(t): |
| embedding = embedding.to(t) |
| return embedding |
|
|
|
|
| |
| |
| |
|
|
|
|
| def patchify_tokens(latent: Tensor, patch: int) -> tuple[Tensor, int, int]: |
| """[B, C, H, W] -> [B, (H/p)(W/p), C*p*p] plus the token grid size.""" |
| b, c, h, w = latent.shape |
| if h % patch or w % patch: |
| raise ValueError(f"latent {h}x{w} is not divisible by patch {patch}") |
| tokens = rearrange(latent, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) |
| return tokens, h // patch, w // patch |
|
|
|
|
| def unpatchify_tokens(tokens: Tensor, h: int, w: int, channels: int, patch: int) -> Tensor: |
| return rearrange(tokens, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h, w=w, c=channels, ph=patch, pw=patch) |
|
|
|
|
| def gather_sequence(sequence: Tensor, indices: Tensor) -> Tensor: |
| """Gather along the token axis with one index list per batch item.""" |
| return sequence.gather(1, indices[..., None].expand(-1, -1, sequence.shape[-1])) |
|
|
|
|
| def scatter_sequence(base: Tensor, indices: Tensor, values: Tensor) -> Tensor: |
| return base.scatter(1, indices[..., None].expand_as(values), values) |
|
|
|
|
| def foreground_indices(mask: Tensor, token_count: int, label: str) -> Tensor: |
| """Flat foreground token indices [B, K] from a [B, h, w] (or [h, w]) mask. |
| |
| Batch size 1 is the supported training configuration (exactly like the |
| FLUX ports); larger batches must have equal foreground counts. |
| """ |
| if mask.ndim == 2: |
| mask = mask[None] |
| flat = mask.reshape(mask.shape[0], -1).bool() |
| if flat.shape[1] != token_count: |
| raise ValueError(f"{label} mask has {flat.shape[1]} cells but the backbone produced {token_count} tokens") |
| counts = flat.sum(dim=1) |
| if torch.any(counts == 0): |
| raise ValueError(f"{label} contains a sample with no foreground tokens") |
| if not torch.all(counts == counts[0]): |
| raise ValueError(f"{label} foreground counts differ within the batch ({counts.tolist()}); use batch_size=1") |
| return torch.stack([row.nonzero(as_tuple=True)[0] for row in flat], dim=0) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def time_shift(mu: float, sigma: float, t: Tensor) -> Tensor: |
| return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) |
|
|
|
|
| def shifted_schedule(num_steps: int, mu: float) -> list[float]: |
| timesteps = torch.linspace(1, 0, num_steps + 1) |
| timesteps = time_shift(mu, 1.0, timesteps) |
| return timesteps.tolist() |
|
|
|
|
| def qwen21_mu(image_seq_len: int, base_shift: float = 0.5, max_shift: float = 0.9, base_len: int = 256, max_len: int = 8192) -> float: |
| """Qwen-Image 2.1 scheduler mu: 0.5 at 256 tokens, 0.9 at 8192 (linear, extrapolated).""" |
| slope = (max_shift - base_shift) / (max_len - base_len) |
| return slope * image_seq_len + (base_shift - slope * base_len) |
|
|
|
|
| |
| |
| |
|
|
|
|
| @dataclass(frozen=True) |
| class LatentNorm: |
| mean: tuple[float, ...] |
| std: tuple[float, ...] |
|
|
| def process_in(self, latent: Tensor) -> Tensor: |
| mean = torch.tensor(self.mean, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1) |
| std = torch.tensor(self.std, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1) |
| return (latent - mean) / std |
|
|
| def process_out(self, latent: Tensor) -> Tensor: |
| mean = torch.tensor(self.mean, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1) |
| std = torch.tensor(self.std, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1) |
| return latent * std + mean |
|
|
|
|
| |
| WAN21_NORM = LatentNorm( |
| mean=(-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, |
| 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921), |
| std=(2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, |
| 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160), |
| ) |
|
|
| |
| QWEN21_NORM = LatentNorm( |
| mean=(0.5126, 0.7721, -0.0631, 1.3506, -0.7855, -2.1025, -0.3458, 1.3722, |
| 1.8873, -1.7177, -0.6510, 0.2732, 0.7562, -0.6163, -1.0277, 3.8363, |
| 2.0210, 0.0472, 0.9320, 2.0087, 2.4954, -0.1391, -1.4249, 1.8464, |
| -0.5236, 1.2826, 3.7046, -1.3035, 2.7286, -1.4518, -1.9036, -1.9955, |
| -0.0342, -1.0265, -0.7636, 3.0555, 0.0746, -3.0751, -0.1076, 1.7376, |
| -1.0914, -1.9435, -0.2784, -1.3680, 0.4809, -0.4433, 0.3764, 0.5729, |
| -2.0595, 1.0960, -1.3260, -2.0211, -5.0179, 0.5275, 4.0162, 1.8505, |
| 0.3026, 1.9373, 1.4937, 0.2632, 0.5547, -1.7121, -0.1562, 0.0304), |
| std=(3.2001, 3.2936, 3.4321, 3.0091, 3.1061, 4.0379, 4.0705, 3.7910, |
| 3.0785, 3.6500, 3.9308, 3.0904, 2.8778, 3.7675, 3.7320, 5.0756, |
| 3.2864, 4.0397, 3.1317, 4.0443, 2.9249, 3.9454, 3.0988, 4.2489, |
| 3.4896, 3.8513, 3.9323, 3.4719, 3.7498, 4.2830, 3.5694, 4.2467, |
| 3.9037, 3.2947, 5.0770, 3.5075, 3.2700, 3.4767, 2.8063, 5.1125, |
| 3.5327, 4.7833, 3.1286, 4.1819, 3.8527, 3.8312, 3.5605, 4.3875, |
| 3.9624, 4.0168, 3.5643, 4.0550, 5.5614, 4.2963, 4.4080, 3.4959, |
| 3.8747, 3.7608, 3.5735, 3.1490, 3.7662, 3.6746, 3.4563, 3.8161), |
| ) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def import_comfy(comfy_root: str | None = None): |
| """Make ComfyUI importable as a library (its CLI parser must not see our argv).""" |
| root = comfy_root or os.environ.get("COMFY_ROOT") |
| if not root or not os.path.isdir(os.path.join(root, "comfy")): |
| raise FileNotFoundError( |
| "COMFY_ROOT must point at a ComfyUI checkout with Qwen-Image 2.1 support " |
| f"(got {root!r})" |
| ) |
| if root not in sys.path: |
| sys.path.insert(0, root) |
| saved_argv = sys.argv |
| sys.argv = [saved_argv[0]] |
| try: |
| import comfy.sd |
| import comfy.utils |
| finally: |
| sys.argv = saved_argv |
| import comfy |
|
|
| return comfy |
|
|
|
|
| class ComfyVAE(nn.Module): |
| """Thin wrapper over ComfyUI's Wan-family VAEs. |
| |
| ``encode``/``decode`` work on single images ``[B, C, H, W]`` in ``[-1, 1]`` |
| and return/consume *normalised* latents (the DiT's input space). The raw |
| autoencoder is ``self.model`` (``comfy.ldm.wan.vae*.WanVAE``); its |
| ``decoder`` + ``conv2`` are what the foreground-aware fine-tune trains. |
| """ |
|
|
| def __init__(self, vae_path: str, norm: LatentNorm, image_channels: int, comfy_root: str | None = None, device="cpu", dtype=torch.float32): |
| super().__init__() |
| comfy = import_comfy(comfy_root) |
| sd = comfy.utils.load_torch_file(vae_path) |
| wrapper = comfy.sd.VAE(sd=sd, device=torch.device("cpu"), dtype=dtype) |
| if wrapper.latent_channels != len(norm.mean): |
| raise ValueError( |
| f"{vae_path} has {wrapper.latent_channels} latent channels, expected {len(norm.mean)}" |
| ) |
| if wrapper.output_channels != image_channels: |
| raise ValueError( |
| f"{vae_path} reconstructs {wrapper.output_channels} image channels, expected {image_channels}" |
| ) |
| self.model = wrapper.first_stage_model.to(device=device, dtype=dtype) |
| self.norm = norm |
| self.image_channels = image_channels |
| self.latent_channels = wrapper.latent_channels |
| self.downscale = wrapper.downscale_ratio if isinstance(wrapper.downscale_ratio, int) else wrapper.downscale_ratio[-1] |
| |
| |
| self.video_only = type(self.model).__module__.endswith("wan.vae") |
|
|
| @property |
| def device(self): |
| return next(self.model.parameters()).device |
|
|
| @property |
| def dtype(self): |
| return next(self.model.parameters()).dtype |
|
|
| def trainable_decoder_parameters(self): |
| return list(self.model.decoder.parameters()) + list(self.model.conv2.parameters()) |
|
|
| def decoder_state_dict(self) -> dict[str, Tensor]: |
| out = {f"decoder.{k}": v for k, v in self.model.decoder.state_dict().items()} |
| out.update({f"conv2.{k}": v for k, v in self.model.conv2.state_dict().items()}) |
| return out |
|
|
| def load_decoder_state_dict(self, state: dict[str, Tensor]) -> None: |
| dec = {k[len("decoder."):]: v for k, v in state.items() if k.startswith("decoder.")} |
| conv2 = {k[len("conv2."):]: v for k, v in state.items() if k.startswith("conv2.")} |
| self.model.decoder.load_state_dict(dec, strict=True) |
| if conv2: |
| self.model.conv2.load_state_dict(conv2, strict=True) |
|
|
| def encode(self, pixels: Tensor) -> Tensor: |
| """pixels [B, C, H, W] in [-1, 1] (C == image_channels) -> normalised latent.""" |
| if pixels.shape[1] != self.image_channels: |
| raise ValueError(f"expected {self.image_channels} image channels, got {pixels.shape[1]}") |
| pixels = pixels.to(self.device, self.dtype) |
| if self.video_only: |
| z = self.model.encode(pixels.unsqueeze(2)).squeeze(2) |
| else: |
| z = self.model.encode(pixels) |
| return self.norm.process_in(z.float()) |
|
|
| def decode(self, latent: Tensor) -> Tensor: |
| """normalised latent [B, C, h, w] -> pixels [B, image_channels, H, W] in [-1, 1].""" |
| z = self.norm.process_out(latent.float()).to(self.dtype) |
| if self.video_only: |
| return self.model.decode(z.unsqueeze(2)).squeeze(2).float() |
| return self.model.decode(z).float() |
|
|
