Add support for Mage-Flow and Mage-Flow Edit
This commit is contained in:
Jaret Burkett
73
extensions_built_in/diffusion_models/mageflow/src/attn.py
Normal file
73
extensions_built_in/diffusion_models/mageflow/src/attn.py
Normal file
@@ -0,0 +1,73 @@
|
||||
"""Varlen attention shim for Mage-Flow.
|
||||
|
||||
The reference implementation (github.com/microsoft/Mage) runs every attention
|
||||
call through flash-attn's ``flash_attn_varlen_func`` so that several
|
||||
variable-length samples can be packed into one sequence and kept isolated via
|
||||
``cu_seqlens``. This shim exposes the same function with the FA2 calling
|
||||
convention, falling back to a per-sequence ``torch.scaled_dot_product_attention``
|
||||
loop when flash-attn is not installed (functionally equivalent for the dense /
|
||||
non-causal path the DiT uses, just slower).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
_RESOLVED_FN: Callable[..., Any] | None = None
|
||||
|
||||
|
||||
def _sdpa_varlen(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q=None,
|
||||
cu_seqlens_k=None,
|
||||
max_seqlen_q=None,
|
||||
max_seqlen_k=None,
|
||||
dropout_p: float = 0.0,
|
||||
softmax_scale=None,
|
||||
causal: bool = False,
|
||||
**_unused: Any,
|
||||
):
|
||||
"""FA2-varlen-compatible SDPA fallback: one SDPA dispatch per sequence."""
|
||||
cu_q = cu_seqlens_q.tolist()
|
||||
cu_k = cu_seqlens_k.tolist()
|
||||
outs = []
|
||||
for qs, qe, ks, ke in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]):
|
||||
# (s, h, d) -> (1, h, s, d)
|
||||
q_i = q[qs:qe].transpose(0, 1).unsqueeze(0)
|
||||
k_i = k[ks:ke].transpose(0, 1).unsqueeze(0)
|
||||
v_i = v[ks:ke].transpose(0, 1).unsqueeze(0)
|
||||
out_i = F.scaled_dot_product_attention(
|
||||
q_i,
|
||||
k_i,
|
||||
v_i,
|
||||
attn_mask=None,
|
||||
dropout_p=dropout_p,
|
||||
is_causal=causal,
|
||||
scale=softmax_scale,
|
||||
)
|
||||
outs.append(out_i.squeeze(0).transpose(0, 1))
|
||||
return torch.cat(outs, dim=0).contiguous()
|
||||
|
||||
|
||||
def _resolve() -> Callable[..., Any]:
|
||||
global _RESOLVED_FN
|
||||
if _RESOLVED_FN is None:
|
||||
try:
|
||||
from flash_attn import flash_attn_varlen_func as _fn
|
||||
|
||||
_RESOLVED_FN = _fn
|
||||
except ImportError:
|
||||
_RESOLVED_FN = _sdpa_varlen
|
||||
return _RESOLVED_FN
|
||||
|
||||
|
||||
def flash_attn_varlen_func(*args, **kwargs):
|
||||
return _resolve()(*args, **kwargs)
|
||||
|
||||
|
||||
__all__ = ["flash_attn_varlen_func"]
|
||||
229
extensions_built_in/diffusion_models/mageflow/src/pipeline.py
Normal file
229
extensions_built_in/diffusion_models/mageflow/src/pipeline.py
Normal file
@@ -0,0 +1,229 @@
|
||||
"""Packing / sampling helpers for Mage-Flow.
|
||||
|
||||
Turns latents + Qwen3-VL text features into the packed variable-length
|
||||
sequences the ``MageFlow`` DiT consumes (batch dim 1, per-sample
|
||||
``cu_seqlens``, exactly mirroring the reference training/inference packing),
|
||||
and provides a minimal flow-matching sampler for ai-toolkit preview images.
|
||||
|
||||
Time convention: rectified flow with ``x_t = (1 - sigma) * clean + sigma *
|
||||
noise``, the model predicts the velocity ``noise - clean``, and sampling
|
||||
integrates sigma from 1 (pure noise) down to 0 (clean). This is identical to
|
||||
ai-toolkit's convention, so the toolkit ``timestep / 1000`` flows straight
|
||||
through as sigma. The inference sigma schedule is the reference's static
|
||||
shift: ``shifted = shift * s / (1 + (shift - 1) * s)`` over ``s =
|
||||
linspace(1, 1/steps, steps)`` with a terminal 0 (shift 6.0).
|
||||
|
||||
Edit conditioning: each sample's sequence is ``[target, ref_1, …, ref_N]`` —
|
||||
clean reference latents ride along after the noisy target tokens, sharing the
|
||||
sample's timestep modulation. RoPE gives every image segment in the pack its
|
||||
own "frame" index (segment order), so references land on later frame
|
||||
coordinates than their target, as in the reference ``generate_edits``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from .transformer import MageFlow
|
||||
|
||||
|
||||
def lens_to_cu(lens: List[int], device) -> torch.Tensor:
|
||||
"""Sequence lengths -> cumulative cu_seqlens [0, l0, l0+l1, ...] (int32)."""
|
||||
t = torch.tensor(lens, device=device, dtype=torch.int32)
|
||||
return torch.cat(
|
||||
[
|
||||
torch.zeros(1, dtype=torch.int32, device=device),
|
||||
torch.cumsum(t, dim=0, dtype=torch.int32),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def pack_text_features(
|
||||
features_list: List[torch.Tensor], device, dtype
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Concatenate per-sample ``(L_i, D)`` text features into the packed
|
||||
``[1, sum L, D]`` text stream + its ``cu_seqlens``. No padding — varlen
|
||||
attention isolates the samples."""
|
||||
feats = [f.to(device, dtype) for f in features_list]
|
||||
txt = torch.cat(feats, dim=0).unsqueeze(0)
|
||||
txt_cu = lens_to_cu([f.shape[0] for f in feats], device)
|
||||
return txt, txt_cu
|
||||
|
||||
|
||||
def predict_velocity(
|
||||
model: MageFlow,
|
||||
latents: torch.Tensor, # (B, C, h, w) noisy target latents
|
||||
t: torch.Tensor, # (B,) flow sigma in [0, 1] (1 = pure noise)
|
||||
text_embeds: List[torch.Tensor], # per-sample (L_i, D) Qwen3-VL features
|
||||
ref_latents: Optional[
|
||||
List[List[torch.Tensor]]
|
||||
] = None, # per-sample clean (C, hr, wr) refs
|
||||
) -> torch.Tensor:
|
||||
"""Run the DiT on the packed [text | target(+refs)] sequences.
|
||||
|
||||
Packs the batch into one varlen forward (mirroring reference packing:
|
||||
per-sample cu_seqlens isolate the samples inside the attention kernel).
|
||||
Reference latents, when given, are appended clean after each sample's
|
||||
target tokens; the returned velocity covers only the target tokens,
|
||||
reshaped back to ``(B, C, h, w)``. No time flip / negation: Mage-Flow's
|
||||
convention matches ai-toolkit's.
|
||||
"""
|
||||
device = latents.device
|
||||
dtype = latents.dtype
|
||||
b, c, h, w = latents.shape
|
||||
|
||||
if ref_latents is not None and not any(len(r) > 0 for r in ref_latents):
|
||||
ref_latents = None
|
||||
|
||||
img_parts, samp_lens, shape_seq, target_idx_parts = [], [], [], []
|
||||
off = 0
|
||||
for i in range(b):
|
||||
tgt = rearrange(latents[i], "c h w -> (h w) c")
|
||||
img_parts.append(tgt)
|
||||
shape_seq.append((1, h, w))
|
||||
target_idx_parts.append(torch.arange(off, off + h * w, device=device))
|
||||
samp_len = h * w
|
||||
off += h * w
|
||||
if ref_latents is not None:
|
||||
for ref in ref_latents[i]:
|
||||
ref = ref.to(device, dtype)
|
||||
_, rh, rw = ref.shape
|
||||
img_parts.append(rearrange(ref, "c h w -> (h w) c"))
|
||||
shape_seq.append((1, rh, rw))
|
||||
samp_len += rh * rw
|
||||
off += rh * rw
|
||||
samp_lens.append(samp_len)
|
||||
|
||||
img = torch.cat(img_parts, dim=0).unsqueeze(0) # [1, sum, C]
|
||||
img_cu = lens_to_cu(samp_lens, device)
|
||||
img_shapes = [shape_seq]
|
||||
target_idx = torch.cat(target_idx_parts)
|
||||
|
||||
txt, txt_cu = pack_text_features(text_embeds, device, dtype)
|
||||
|
||||
out = model(
|
||||
img=img,
|
||||
txt=txt,
|
||||
timesteps=t.to(device),
|
||||
img_shapes=img_shapes,
|
||||
img_cu_seqlens=img_cu,
|
||||
txt_cu_seqlens=txt_cu,
|
||||
) # [1, sum, C_out]
|
||||
|
||||
pred = out[:, target_idx, :] # [1, B*h*w, C]
|
||||
return rearrange(pred, "1 (b h w) c -> b c h w", b=b, h=h, w=w)
|
||||
|
||||
|
||||
def build_shifted_sigmas(
|
||||
num_steps: int, shift: float = 6.0, device=None
|
||||
) -> torch.Tensor:
|
||||
"""The reference inference sigma schedule: ``linspace(1, 1/steps, steps)``
|
||||
run through the static shift ``shift*s/(1+(shift-1)*s)`` with a terminal 0."""
|
||||
s = torch.linspace(1.0, 1.0 / num_steps, num_steps, dtype=torch.float64)
|
||||
s = shift * s / (1 + (shift - 1) * s)
|
||||
sigmas = torch.cat([s, torch.zeros(1, dtype=torch.float64)])
|
||||
return sigmas.to(device=device, dtype=torch.float32)
|
||||
|
||||
|
||||
class MageFlowPipeline:
|
||||
"""Lightweight flow-matching sampler used by ai-toolkit's preview generation."""
|
||||
|
||||
def __init__(self, model):
|
||||
# ``model`` is the MageFlowModel (BaseModel) so we can reuse its
|
||||
# encode/decode and config.
|
||||
self.model = model
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.model.device_torch
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def set_progress_bar_config(self, **kwargs):
|
||||
pass
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
conditional_embeds,
|
||||
unconditional_embeds,
|
||||
height: int = 1024,
|
||||
width: int = 1024,
|
||||
num_inference_steps: int = 30,
|
||||
guidance_scale: float = 5.0,
|
||||
latents: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
ref_latents: Optional[List[List[torch.Tensor]]] = None,
|
||||
**kwargs,
|
||||
) -> List[Image.Image]:
|
||||
model = self.model
|
||||
device = model.device_torch
|
||||
dtype = model.torch_dtype
|
||||
transformer: MageFlow = model.transformer
|
||||
|
||||
shift = float(model.model_config.model_kwargs.get("static_shift", 6.0))
|
||||
renorm = bool(model.model_config.model_kwargs.get("cfg_renormalization", False))
|
||||
|
||||
do_cfg = guidance_scale > 1.0 and unconditional_embeds is not None
|
||||
|
||||
gh = height // model.vae_scale_factor
|
||||
gw = width // model.vae_scale_factor
|
||||
latent_channels = transformer.in_channels
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(
|
||||
(1, latent_channels, gh, gw),
|
||||
generator=generator,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
latents = latents.to(device, dtype=torch.float32)
|
||||
|
||||
cond_feats = [f.to(device, dtype) for f in conditional_embeds.text_embeds]
|
||||
if do_cfg:
|
||||
uncond_feats = [
|
||||
f.to(device, dtype) for f in unconditional_embeds.text_embeds
|
||||
]
|
||||
|
||||
sigmas = build_shifted_sigmas(num_inference_steps, shift=shift, device=device)
|
||||
|
||||
# Euler integration of the flow ODE, sigma 1 -> 0 (with optional CFG:
|
||||
# v = uncond + cfg * (cond - uncond), reference convention).
|
||||
for i in range(num_inference_steps):
|
||||
s_cur = sigmas[i].item()
|
||||
s_next = sigmas[i + 1].item()
|
||||
t = torch.full((latents.shape[0],), s_cur, dtype=dtype, device=device)
|
||||
v_cond = predict_velocity(
|
||||
transformer, latents.to(dtype), t, cond_feats, ref_latents=ref_latents
|
||||
)
|
||||
if do_cfg:
|
||||
v_uncond = predict_velocity(
|
||||
transformer,
|
||||
latents.to(dtype),
|
||||
t,
|
||||
uncond_feats,
|
||||
ref_latents=ref_latents,
|
||||
)
|
||||
v = v_uncond + guidance_scale * (v_cond - v_uncond)
|
||||
if renorm:
|
||||
# Rescale the guided velocity per token back to the
|
||||
# conditional velocity's norm (reduces oversaturation).
|
||||
v = v * (
|
||||
torch.norm(v_cond, dim=1, keepdim=True)
|
||||
/ (torch.norm(v, dim=1, keepdim=True) + 1e-6)
|
||||
)
|
||||
else:
|
||||
v = v_cond
|
||||
latents = latents + (s_next - s_cur) * v.to(torch.float32)
|
||||
|
||||
images = model.decode_latents(latents, device=device, dtype=dtype)
|
||||
images = images.float().clamp(-1.0, 1.0)
|
||||
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
|
||||
images = images.permute(0, 2, 3, 1).cpu().numpy()
|
||||
return [Image.fromarray(arr) for arr in images]
|
||||
@@ -0,0 +1,169 @@
|
||||
"""Qwen3-VL prompt encoding for Mage-Flow.
|
||||
|
||||
The reference implementation (mage_flow/models/modules/text_encoder.py) packs
|
||||
many prompts into one varlen forward; per-sequence the math reduces to a plain
|
||||
causal forward with positions ``arange(L)`` replicated on every mrope axis and
|
||||
the leading system-prompt tokens dropped from the output. ai-toolkit encodes
|
||||
prompts one at a time (embeds are cached per caption), so that per-sequence
|
||||
form is implemented here directly on the stock HF Qwen3-VL module.
|
||||
|
||||
The DiT consumes the final hidden states ``[L - drop_idx, 2560]``; the pooled
|
||||
"vec" of the reference is unused by the released checkpoints (vec_type null —
|
||||
the transformer adds a zero vector), so it is not computed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx``
|
||||
# is the number of leading (system prompt) tokens dropped from the encoder
|
||||
# output before it conditions the DiT.
|
||||
PROMPT_TEMPLATE = {
|
||||
"mage-flow": {
|
||||
"template": (
|
||||
"<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, "
|
||||
"text, spatial relationships of the objects and background:"
|
||||
"<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
|
||||
),
|
||||
"start_idx": 34,
|
||||
},
|
||||
"mage-flow-edit": {
|
||||
"template": (
|
||||
"<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture,"
|
||||
" objects, background), then explain how the user's text instruction should alter or modify the image. "
|
||||
"Generate a new image that meets the user's requirements while maintaining consistency with the original "
|
||||
"input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
|
||||
),
|
||||
"start_idx": 64,
|
||||
},
|
||||
}
|
||||
|
||||
# Fixed image placeholder used at edit training time (one per reference image).
|
||||
EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
|
||||
|
||||
def edit_prompt_body(instruction: str, num_refs: int) -> str:
|
||||
"""Training-time multi-reference prompt body:
|
||||
``Image 1: <ph>Image 2: <ph>…{instruction}``."""
|
||||
prefix = "".join(
|
||||
f"Image {j}: {EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1)
|
||||
)
|
||||
return prefix + instruction
|
||||
|
||||
|
||||
def patch_qwen_vl_patch_embed(model):
|
||||
"""Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain
|
||||
linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
|
||||
falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear
|
||||
(a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves.
|
||||
Returns the number of patch_embed modules patched. (Same patch as the krea2
|
||||
extension / Qwen3VLCaptioner.)"""
|
||||
patched = 0
|
||||
for module in model.modules():
|
||||
proj = getattr(module, "proj", None)
|
||||
if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
|
||||
proj.stride
|
||||
):
|
||||
|
||||
def fast_forward(hidden_states, _proj=proj):
|
||||
w = _proj.weight.reshape(_proj.weight.shape[0], -1)
|
||||
x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
|
||||
return F.linear(x, w, _proj.bias)
|
||||
|
||||
module.forward = fast_forward
|
||||
patched += 1
|
||||
return patched
|
||||
|
||||
|
||||
def resize_vl_images(
|
||||
images: List[torch.Tensor], max_long_edge: int = 384
|
||||
) -> List["PIL.Image.Image"]:
|
||||
"""Prepare reference images for the Qwen3-VL conditioning pass.
|
||||
|
||||
Matches the reference's ``_resize_long_edge`` (cap the long edge at
|
||||
``max_long_edge``, preserving aspect ratio, BICUBIC, never upscale) — the
|
||||
MLLM only needs a coarse view of the reference; full-resolution detail
|
||||
flows through the VAE reference latents. Input tensors are ``(C, H, W)``
|
||||
or ``(1, C, H, W)`` in [0, 1]; output is PIL for the Qwen processor.
|
||||
"""
|
||||
from torchvision.transforms.functional import to_pil_image
|
||||
|
||||
out = []
|
||||
for img in images:
|
||||
if img.dim() == 4:
|
||||
img = img[0]
|
||||
img = img.float().clamp(0, 1).cpu()
|
||||
pil = to_pil_image(img)
|
||||
if max_long_edge and max_long_edge > 0:
|
||||
w, h = pil.size
|
||||
long_edge = max(w, h)
|
||||
if long_edge > max_long_edge:
|
||||
scale = max_long_edge / long_edge
|
||||
new_w = max(1, int(round(w * scale)))
|
||||
new_h = max(1, int(round(h * scale)))
|
||||
pil = pil.resize((new_w, new_h), resample=3) # PIL.Image.BICUBIC
|
||||
out.append(pil)
|
||||
return out
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_mageflow_prompt(
|
||||
text_encoder, # Qwen3VLForConditionalGeneration
|
||||
tokenizer,
|
||||
prompt: str,
|
||||
template_name: str = "mage-flow",
|
||||
max_length: int = 2048,
|
||||
images: Optional[list] = None, # list of PIL images (already VL-resized)
|
||||
processor=None, # AutoProcessor, required when images are given
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> torch.Tensor:
|
||||
"""Encode one prompt (optionally with reference images) to DiT conditioning.
|
||||
|
||||
Returns ``(L - start_idx, 2560)`` final hidden states with the system
|
||||
prompt dropped, matching the reference ``TextEncoder.forward`` txt output.
|
||||
"""
|
||||
info = PROMPT_TEMPLATE[template_name]
|
||||
template = info["template"]
|
||||
drop_idx = int(info["start_idx"])
|
||||
device = next(text_encoder.parameters()).device
|
||||
|
||||
if images:
|
||||
if processor is None:
|
||||
raise ValueError(
|
||||
"encoding reference images requires the Qwen3-VL AutoProcessor"
|
||||
)
|
||||
formatted = template.format(edit_prompt_body(prompt, len(images)))
|
||||
vl = processor(text=[formatted], images=list(images), return_tensors="pt")
|
||||
input_ids = vl["input_ids"].to(device)
|
||||
extra = {}
|
||||
for key in ("pixel_values", "image_grid_thw"):
|
||||
if vl.get(key) is not None:
|
||||
extra[key] = vl[key].to(device)
|
||||
else:
|
||||
formatted = template.format(prompt)
|
||||
input_ids = tokenizer(
|
||||
formatted,
|
||||
max_length=max_length + drop_idx,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
).input_ids.to(device)
|
||||
extra = {}
|
||||
|
||||
# The reference packed encoder gives every sequence positions arange(L)
|
||||
# replicated on all mrope axes (even for image tokens); a 2D position_ids
|
||||
# is expanded to exactly that by the HF text model.
|
||||
position_ids = torch.arange(input_ids.shape[1], device=device).unsqueeze(0)
|
||||
|
||||
out = text_encoder.model(
|
||||
input_ids=input_ids,
|
||||
position_ids=position_ids,
|
||||
**extra,
|
||||
)
|
||||
hidden = out.last_hidden_state[0] # (L, D)
|
||||
return hidden[drop_idx:].to(dtype)
|
||||
809
extensions_built_in/diffusion_models/mageflow/src/transformer.py
Normal file
809
extensions_built_in/diffusion_models/mageflow/src/transformer.py
Normal file
@@ -0,0 +1,809 @@
|
||||
"""Mage-Flow NR-MMDiT (dual-stream, native-resolution packed) for ai-toolkit.
|
||||
|
||||
Vendored from the reference implementation
|
||||
(github.com/microsoft/Mage, mage_flow/models/mage_flow.py + modules/mage_layers.py)
|
||||
with these deviations:
|
||||
|
||||
- loguru / pydantic / torch._dynamo decorators dropped,
|
||||
- gradient checkpointing is gated on ``self.gradient_checkpointing and
|
||||
torch.is_grad_enabled()`` (instead of the reference's ``self.training and
|
||||
self.checkpoint``) and exposed via ``enable_gradient_checkpointing`` /
|
||||
``disable_gradient_checkpointing`` so ai-toolkit's trainer can toggle it,
|
||||
- flash-attn is imported through the local shim (``.attn``) which falls back
|
||||
to a per-sequence SDPA loop when flash-attn is unavailable.
|
||||
|
||||
State-dict keys are identical to the reference checkpoints
|
||||
(``transformer/diffusion_pytorch_model.safetensors`` in
|
||||
microsoft/Mage-Flow-Base and microsoft/Mage-Flow-Edit-Base).
|
||||
|
||||
Sequence layout: images and text are packed as variable-length sequences of
|
||||
batch dim 1 (``[1, sum_len, C]``) with per-sample ``cu_seqlens``; joint text+image
|
||||
attention runs in a single varlen kernel. 2D multi-scale RoPE is applied to
|
||||
image tokens only (text tokens are not rotated). ``timesteps`` is the flow
|
||||
sigma in [0, 1] (1 = pure noise), one entry per packed sample.
|
||||
"""
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
|
||||
from diffusers.models.attention import FeedForward
|
||||
from diffusers.models.embeddings import TimestepEmbedding
|
||||
from diffusers.models.normalization import RMSNorm
|
||||
|
||||
from .attn import flash_attn_varlen_func
|
||||
|
||||
|
||||
def apply_rotary_emb_mageflow(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply complex rotary embeddings to `x` ([S, H, D]) using `freqs_cis`
|
||||
(the MageFlowEmbedRope 2D multi-scale RoPE, adjacent-pair complex convention)."""
|
||||
x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
|
||||
freqs_cis = freqs_cis.unsqueeze(1)
|
||||
x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(-2)
|
||||
return x_out.type_as(x)
|
||||
|
||||
|
||||
def get_timestep_embedding(
|
||||
timesteps: torch.Tensor,
|
||||
embedding_dim: int,
|
||||
flip_sin_to_cos: bool = False,
|
||||
downscale_freq_shift: float = 1,
|
||||
scale: float = 1,
|
||||
max_period: int = 10000,
|
||||
) -> torch.Tensor:
|
||||
"""Sinusoidal timestep embeddings (DDPM convention).
|
||||
|
||||
NOTE: kept vendored (not diffusers') because the frequency table is
|
||||
downcast to ``timesteps.dtype`` (bf16) here — the model was trained with
|
||||
this exact bf16 rounding, so diffusers' fp32 variant produces a slightly
|
||||
different embedding and degrades outputs.
|
||||
"""
|
||||
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
exponent = -math.log(max_period) * torch.arange(
|
||||
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
|
||||
)
|
||||
exponent = exponent / (half_dim - downscale_freq_shift)
|
||||
|
||||
emb = torch.exp(exponent).to(timesteps.dtype)
|
||||
emb = timesteps[:, None].float() * emb[None, :]
|
||||
|
||||
# scale embeddings
|
||||
emb = scale * emb
|
||||
|
||||
# concat sine and cosine embeddings
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
|
||||
# flip sine and cosine embeddings
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
|
||||
# zero pad
|
||||
if embedding_dim % 2 == 1:
|
||||
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
|
||||
return emb
|
||||
|
||||
|
||||
class Timesteps(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
num_channels: int,
|
||||
flip_sin_to_cos: bool,
|
||||
downscale_freq_shift: float,
|
||||
scale: int = 1,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_channels = num_channels
|
||||
self.flip_sin_to_cos = flip_sin_to_cos
|
||||
self.downscale_freq_shift = downscale_freq_shift
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
|
||||
return get_timestep_embedding(
|
||||
timesteps,
|
||||
self.num_channels,
|
||||
flip_sin_to_cos=self.flip_sin_to_cos,
|
||||
downscale_freq_shift=self.downscale_freq_shift,
|
||||
scale=self.scale,
|
||||
)
|
||||
|
||||
|
||||
class MageFlowTimestepProjEmbeddings(nn.Module):
|
||||
def __init__(self, embedding_dim):
|
||||
super().__init__()
|
||||
self.time_proj = Timesteps(
|
||||
num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000
|
||||
)
|
||||
self.timestep_embedder = TimestepEmbedding(
|
||||
in_channels=256, time_embed_dim=embedding_dim
|
||||
)
|
||||
|
||||
def forward(self, timestep, hidden_states):
|
||||
timesteps_proj = self.time_proj(timestep)
|
||||
timesteps_emb = self.timestep_embedder(
|
||||
timesteps_proj.to(dtype=hidden_states.dtype)
|
||||
) # (N, D)
|
||||
return timesteps_emb
|
||||
|
||||
|
||||
class MageFlowEmbedRope(nn.Module):
|
||||
def __init__(self, theta: int, axes_dim: list, scale_rope=False):
|
||||
super().__init__()
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
pos_index = torch.arange(4096)
|
||||
neg_index = torch.arange(4096).flip(0) * -1 - 1
|
||||
self.pos_freqs = torch.cat(
|
||||
[
|
||||
self.rope_params(pos_index, self.axes_dim[0], self.theta),
|
||||
self.rope_params(pos_index, self.axes_dim[1], self.theta),
|
||||
self.rope_params(pos_index, self.axes_dim[2], self.theta),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
self.neg_freqs = torch.cat(
|
||||
[
|
||||
self.rope_params(neg_index, self.axes_dim[0], self.theta),
|
||||
self.rope_params(neg_index, self.axes_dim[1], self.theta),
|
||||
self.rope_params(neg_index, self.axes_dim[2], self.theta),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
|
||||
self.scale_rope = scale_rope
|
||||
self.video_freq_cache = {}
|
||||
|
||||
def rope_params(self, index, dim, theta=10000):
|
||||
assert dim % 2 == 0
|
||||
freqs = torch.outer(
|
||||
index,
|
||||
1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)),
|
||||
)
|
||||
freqs = torch.polar(torch.ones_like(freqs), freqs)
|
||||
return freqs
|
||||
|
||||
def forward(
|
||||
self, video_fhw, device: torch.device, max_img_len: int = None
|
||||
) -> torch.Tensor:
|
||||
"""Compute the vision RoPE frequencies for the packed image tokens.
|
||||
Text tokens are NOT rotated, so no text RoPE is computed.
|
||||
|
||||
``video_fhw`` is a list of (frame, height, width) tuples, one per image
|
||||
segment in the packed sequence; segment index doubles as the "frame"
|
||||
RoPE coordinate (edit reference images land on later frame indices).
|
||||
"""
|
||||
if self.pos_freqs.device != device:
|
||||
self.pos_freqs = self.pos_freqs.to(device)
|
||||
self.neg_freqs = self.neg_freqs.to(device)
|
||||
|
||||
if isinstance(video_fhw, list):
|
||||
video_fhw = video_fhw[0]
|
||||
if not isinstance(video_fhw, list):
|
||||
video_fhw = [video_fhw]
|
||||
|
||||
vid_freqs = []
|
||||
for idx, fhw in enumerate(video_fhw):
|
||||
frame, height, width = fhw
|
||||
key = (frame, height, width, idx)
|
||||
if key not in self.video_freq_cache:
|
||||
self.video_freq_cache[key] = self._compute_video_freqs(
|
||||
frame, height, width, idx
|
||||
)
|
||||
vid_freqs.append(self.video_freq_cache[key].to(device))
|
||||
|
||||
vid_freqs = torch.cat(vid_freqs, dim=0)
|
||||
|
||||
if max_img_len is not None and vid_freqs.shape[0] < max_img_len:
|
||||
pad_len = max_img_len - vid_freqs.shape[0]
|
||||
vid_freqs = torch.nn.functional.pad(vid_freqs, (0, 0, 0, pad_len))
|
||||
|
||||
return vid_freqs
|
||||
|
||||
def _compute_video_freqs(
|
||||
self, frame: int, height: int, width: int, idx: int = 0
|
||||
) -> torch.Tensor:
|
||||
seq_lens = frame * height * width
|
||||
freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
|
||||
freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
|
||||
|
||||
freqs_frame = (
|
||||
freqs_pos[0][idx : idx + frame]
|
||||
.view(frame, 1, 1, -1)
|
||||
.expand(frame, height, width, -1)
|
||||
)
|
||||
if self.scale_rope:
|
||||
freqs_height = torch.cat(
|
||||
[freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
|
||||
dim=0,
|
||||
)
|
||||
freqs_height = freqs_height.view(1, height, 1, -1).expand(
|
||||
frame, height, width, -1
|
||||
)
|
||||
freqs_width = torch.cat(
|
||||
[freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
|
||||
dim=0,
|
||||
)
|
||||
freqs_width = freqs_width.view(1, 1, width, -1).expand(
|
||||
frame, height, width, -1
|
||||
)
|
||||
else:
|
||||
freqs_height = (
|
||||
freqs_pos[1][:height]
|
||||
.view(1, height, 1, -1)
|
||||
.expand(frame, height, width, -1)
|
||||
)
|
||||
freqs_width = (
|
||||
freqs_pos[2][:width]
|
||||
.view(1, 1, width, -1)
|
||||
.expand(frame, height, width, -1)
|
||||
)
|
||||
|
||||
freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(
|
||||
seq_lens, -1
|
||||
)
|
||||
return freqs.clone().contiguous()
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""Joint text+image attention (reference subset: the double-stream config
|
||||
MageFlow actually instantiates — qk rms_norm, added kv projections)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
query_dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = False,
|
||||
added_kv_proj_dim: int = None,
|
||||
added_proj_bias: bool = True,
|
||||
out_bias: bool = True,
|
||||
eps: float = 1e-5,
|
||||
processor=None,
|
||||
out_dim: int = None,
|
||||
out_context_dim: int = None,
|
||||
elementwise_affine: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
|
||||
self.inner_kv_dim = self.inner_dim
|
||||
self.query_dim = query_dim
|
||||
self.use_bias = bias
|
||||
self.cross_attention_dim = query_dim
|
||||
self.out_dim = out_dim if out_dim is not None else query_dim
|
||||
self.out_context_dim = (
|
||||
out_context_dim if out_context_dim is not None else query_dim
|
||||
)
|
||||
self.heads = out_dim // dim_head if out_dim is not None else heads
|
||||
self.added_kv_proj_dim = added_kv_proj_dim
|
||||
self.added_proj_bias = added_proj_bias
|
||||
|
||||
# qk_norm is always "rms_norm" for MageFlow.
|
||||
self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
|
||||
self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
|
||||
|
||||
self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
|
||||
self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
|
||||
self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
|
||||
|
||||
self.add_k_proj = nn.Linear(
|
||||
added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias
|
||||
)
|
||||
self.add_v_proj = nn.Linear(
|
||||
added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias
|
||||
)
|
||||
self.add_q_proj = nn.Linear(
|
||||
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
|
||||
)
|
||||
self.norm_added_q = RMSNorm(dim_head, eps=eps)
|
||||
self.norm_added_k = RMSNorm(dim_head, eps=eps)
|
||||
|
||||
self.to_out = nn.ModuleList([])
|
||||
self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
|
||||
self.to_out.append(nn.Dropout(dropout))
|
||||
|
||||
self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
|
||||
|
||||
self.processor = processor
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
txt_cu_lens: torch.Tensor = None,
|
||||
img_cu_lens: torch.Tensor = None,
|
||||
image_rotary_emb: torch.Tensor = None,
|
||||
**attention_kwargs,
|
||||
) -> torch.Tensor:
|
||||
return self.processor(
|
||||
self,
|
||||
hidden_states,
|
||||
txt_cu_lens=txt_cu_lens,
|
||||
img_cu_lens=img_cu_lens,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
**attention_kwargs,
|
||||
)
|
||||
|
||||
|
||||
class MageDoubleStreamAttnProcessor:
|
||||
"""Attention processor for the Mage double-stream architecture. Implements
|
||||
joint attention where the packed text and image streams are concatenated
|
||||
per sample (order: [text, image]) and run through one varlen kernel."""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states: torch.FloatTensor, # Image stream [1, sum_img, D]
|
||||
img_cu_lens: torch.LongTensor,
|
||||
encoder_hidden_states: torch.FloatTensor = None, # Text stream [1, sum_txt, D]
|
||||
txt_cu_lens: torch.LongTensor = None,
|
||||
image_rotary_emb: torch.Tensor = None,
|
||||
**kwargs,
|
||||
) -> torch.FloatTensor:
|
||||
if encoder_hidden_states is None:
|
||||
raise ValueError(
|
||||
"MageDoubleStreamAttnProcessor requires encoder_hidden_states (text stream)"
|
||||
)
|
||||
|
||||
# Compute QKV for image stream (sample projections)
|
||||
img_query = attn.to_q(hidden_states)
|
||||
img_key = attn.to_k(hidden_states)
|
||||
img_value = attn.to_v(hidden_states)
|
||||
|
||||
# Compute QKV for text stream (context projections)
|
||||
txt_query = attn.add_q_proj(encoder_hidden_states)
|
||||
txt_key = attn.add_k_proj(encoder_hidden_states)
|
||||
txt_value = attn.add_v_proj(encoder_hidden_states)
|
||||
|
||||
# Reshape for multi-head attention
|
||||
img_query = img_query.unflatten(-1, (attn.heads, -1))
|
||||
img_key = img_key.unflatten(-1, (attn.heads, -1))
|
||||
img_value = img_value.unflatten(-1, (attn.heads, -1))
|
||||
|
||||
txt_query = txt_query.unflatten(-1, (attn.heads, -1))
|
||||
txt_key = txt_key.unflatten(-1, (attn.heads, -1))
|
||||
txt_value = txt_value.unflatten(-1, (attn.heads, -1))
|
||||
|
||||
if img_query.ndim == 4:
|
||||
img_query = img_query.flatten(0, 1)
|
||||
img_key = img_key.flatten(0, 1)
|
||||
img_value = img_value.flatten(0, 1)
|
||||
|
||||
if txt_query.ndim == 4:
|
||||
txt_query = txt_query.flatten(0, 1)
|
||||
txt_key = txt_key.flatten(0, 1)
|
||||
txt_value = txt_value.flatten(0, 1)
|
||||
|
||||
# Apply QK normalization
|
||||
img_query = attn.norm_q(img_query)
|
||||
img_key = attn.norm_k(img_key)
|
||||
txt_query = attn.norm_added_q(txt_query)
|
||||
txt_key = attn.norm_added_k(txt_key)
|
||||
|
||||
# Apply 2D multi-scale RoPE (MageFlowEmbedRope) to image tokens
|
||||
img_freqs = image_rotary_emb
|
||||
img_query = apply_rotary_emb_mageflow(img_query, img_freqs)
|
||||
img_key = apply_rotary_emb_mageflow(img_key, img_freqs)
|
||||
|
||||
# Calculate lengths
|
||||
img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
|
||||
txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
|
||||
|
||||
# Calculate joint cu_seqlens
|
||||
joint_lens = txt_lens + img_lens
|
||||
joint_cu_lens = torch.cat(
|
||||
[
|
||||
torch.zeros(1, dtype=torch.int32, device=joint_lens.device),
|
||||
torch.cumsum(joint_lens, dim=0, dtype=torch.int32),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
device = joint_lens.device
|
||||
batch_size = len(txt_lens)
|
||||
sample_indices = torch.arange(batch_size, device=device)
|
||||
|
||||
txt_sample_ids = torch.repeat_interleave(sample_indices, txt_lens)
|
||||
img_sample_ids = torch.repeat_interleave(sample_indices, img_lens)
|
||||
|
||||
txt_intra_pos = (
|
||||
torch.arange(txt_query.shape[0], device=device)
|
||||
- txt_cu_lens[txt_sample_ids]
|
||||
)
|
||||
img_intra_pos = (
|
||||
torch.arange(img_query.shape[0], device=device)
|
||||
- img_cu_lens[img_sample_ids]
|
||||
)
|
||||
|
||||
txt_dest_indices = joint_cu_lens[txt_sample_ids] + txt_intra_pos
|
||||
img_dest_indices = (
|
||||
joint_cu_lens[img_sample_ids] + txt_lens[img_sample_ids] + img_intra_pos
|
||||
)
|
||||
|
||||
total_tokens = joint_cu_lens[-1]
|
||||
joint_query = torch.empty(
|
||||
(total_tokens, *txt_query.shape[1:]), dtype=txt_query.dtype, device=device
|
||||
)
|
||||
joint_key = torch.empty(
|
||||
(total_tokens, *txt_key.shape[1:]), dtype=txt_key.dtype, device=device
|
||||
)
|
||||
joint_value = torch.empty(
|
||||
(total_tokens, *txt_value.shape[1:]), dtype=txt_value.dtype, device=device
|
||||
)
|
||||
|
||||
joint_query[txt_dest_indices] = txt_query
|
||||
joint_query[img_dest_indices] = img_query
|
||||
|
||||
joint_key[txt_dest_indices] = txt_key
|
||||
joint_key[img_dest_indices] = img_key
|
||||
|
||||
joint_value[txt_dest_indices] = txt_value
|
||||
joint_value[img_dest_indices] = img_value
|
||||
|
||||
max_seqlen = joint_lens.max().item()
|
||||
joint_attn_output = flash_attn_varlen_func(
|
||||
joint_query,
|
||||
joint_key,
|
||||
joint_value,
|
||||
cu_seqlens_q=joint_cu_lens,
|
||||
cu_seqlens_k=joint_cu_lens,
|
||||
max_seqlen_q=max_seqlen,
|
||||
max_seqlen_k=max_seqlen,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
txt_attn_output = joint_attn_output[txt_dest_indices]
|
||||
img_attn_output = joint_attn_output[img_dest_indices]
|
||||
|
||||
img_attn_output = img_attn_output.flatten(1, 2) # (N, H, D) -> (N, H*D)
|
||||
img_attn_output = img_attn_output.to(joint_query.dtype)
|
||||
|
||||
txt_attn_output = txt_attn_output.flatten(1, 2) # (N, H, D) -> (N, H*D)
|
||||
txt_attn_output = txt_attn_output.to(joint_query.dtype)
|
||||
|
||||
img_attn_output = attn.to_out[0](img_attn_output)
|
||||
img_attn_output = attn.to_out[1](img_attn_output) # dropout
|
||||
|
||||
txt_attn_output = attn.to_add_out(txt_attn_output)
|
||||
txt_attn_output = txt_attn_output.view(
|
||||
encoder_hidden_states.shape[0],
|
||||
encoder_hidden_states.shape[1],
|
||||
txt_attn_output.shape[-1],
|
||||
)
|
||||
|
||||
return img_attn_output, txt_attn_output
|
||||
|
||||
|
||||
class MageFlowTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
eps: float = 1e-6,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
|
||||
# Image processing modules
|
||||
self.img_mod = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(
|
||||
dim, 6 * dim, bias=True
|
||||
), # For scale, shift, gate for norm1 and norm2
|
||||
)
|
||||
self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
||||
self.attn = Attention(
|
||||
query_dim=dim,
|
||||
added_kv_proj_dim=dim, # Enable added KV projections for text stream
|
||||
dim_head=attention_head_dim,
|
||||
heads=num_attention_heads,
|
||||
out_dim=dim,
|
||||
bias=True,
|
||||
processor=MageDoubleStreamAttnProcessor(),
|
||||
eps=eps,
|
||||
)
|
||||
self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
||||
self.img_mlp = FeedForward(
|
||||
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
|
||||
)
|
||||
|
||||
# Text processing modules
|
||||
self.txt_mod = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(
|
||||
dim, 6 * dim, bias=True
|
||||
), # For scale, shift, gate for norm1 and norm2
|
||||
)
|
||||
self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
||||
# Text doesn't need separate attention - it's handled by img attn joint computation
|
||||
self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
||||
self.txt_mlp = FeedForward(
|
||||
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
|
||||
)
|
||||
|
||||
def _modulate(self, x, mod_params, cu_lens=None):
|
||||
"""Apply modulation to input tensor"""
|
||||
shift, scale, gate = mod_params.chunk(3, dim=-1)
|
||||
if cu_lens is not None:
|
||||
assert x.shape[0] == 1, "x must be of shape (1, *) when cu_lens is not None"
|
||||
x_flattened = x.view(-1, x.shape[-1])
|
||||
lengths = cu_lens[1:] - cu_lens[:-1]
|
||||
shift_t = shift.repeat_interleave(lengths, dim=0)
|
||||
scale_t = scale.repeat_interleave(lengths, dim=0)
|
||||
gate_t = gate.repeat_interleave(lengths, dim=0)
|
||||
|
||||
x_flattened = x_flattened * (1 + scale_t) + shift_t
|
||||
x = x_flattened.view(x.shape)
|
||||
return x, gate_t
|
||||
else:
|
||||
return x * (1 + scale) + shift, gate
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
image_rotary_emb: torch.Tensor,
|
||||
txt_cu_lens: torch.Tensor,
|
||||
img_cu_lens: torch.Tensor,
|
||||
joint_attention_kwargs: dict = None,
|
||||
):
|
||||
# Get modulation parameters for both streams
|
||||
img_mod_params = self.img_mod(temb) # [B, 6*dim]
|
||||
txt_mod_params = self.txt_mod(temb) # [B, 6*dim]
|
||||
|
||||
# Split modulation parameters for norm1 and norm2
|
||||
img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
|
||||
txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
|
||||
|
||||
# Process image stream - norm1 + modulation
|
||||
img_normed = self.img_norm1(hidden_states)
|
||||
img_modulated, img_gate1 = self._modulate(
|
||||
img_normed, img_mod1, cu_lens=img_cu_lens
|
||||
)
|
||||
|
||||
# Process text stream - norm1 + modulation
|
||||
txt_normed = self.txt_norm1(encoder_hidden_states)
|
||||
txt_modulated, txt_gate1 = self._modulate(
|
||||
txt_normed, txt_mod1, cu_lens=txt_cu_lens
|
||||
)
|
||||
|
||||
# Joint attention: computes QKV for both streams, applies QK norm and
|
||||
# RoPE, concatenates per sample and splits the results back.
|
||||
joint_attention_kwargs = joint_attention_kwargs or {}
|
||||
attn_output = self.attn(
|
||||
hidden_states=img_modulated,
|
||||
encoder_hidden_states=txt_modulated,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
txt_cu_lens=txt_cu_lens,
|
||||
img_cu_lens=img_cu_lens,
|
||||
**joint_attention_kwargs,
|
||||
)
|
||||
|
||||
img_attn_output, txt_attn_output = attn_output
|
||||
|
||||
# Apply attention gates and add residual
|
||||
hidden_states = hidden_states + img_gate1 * img_attn_output
|
||||
encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
|
||||
|
||||
# Process image stream - norm2 + MLP
|
||||
img_normed2 = self.img_norm2(hidden_states)
|
||||
img_modulated2, img_gate2 = self._modulate(
|
||||
img_normed2, img_mod2, cu_lens=img_cu_lens
|
||||
)
|
||||
img_mlp_output = self.img_mlp(img_modulated2)
|
||||
hidden_states = hidden_states + img_gate2 * img_mlp_output
|
||||
|
||||
# Process text stream - norm2 + MLP
|
||||
txt_normed2 = self.txt_norm2(encoder_hidden_states)
|
||||
txt_modulated2, txt_gate2 = self._modulate(
|
||||
txt_normed2, txt_mod2, cu_lens=txt_cu_lens
|
||||
)
|
||||
txt_mlp_output = self.txt_mlp(txt_modulated2)
|
||||
encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
|
||||
|
||||
# Clip to prevent overflow for fp16
|
||||
if encoder_hidden_states.dtype == torch.float16:
|
||||
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
|
||||
if hidden_states.dtype == torch.float16:
|
||||
hidden_states = hidden_states.clip(-65504, 65504)
|
||||
|
||||
return encoder_hidden_states, hidden_states
|
||||
|
||||
|
||||
class AdaLayerNormContinuous(nn.Module):
|
||||
"""Adaptive output norm; supports packed sequences via ``cu_seqlens``."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
conditioning_embedding_dim: int,
|
||||
elementwise_affine=True,
|
||||
eps=1e-5,
|
||||
bias=True,
|
||||
norm_type="layer_norm",
|
||||
):
|
||||
super().__init__()
|
||||
self.silu = nn.SiLU()
|
||||
self.linear = nn.Linear(
|
||||
conditioning_embedding_dim, embedding_dim * 2, bias=bias
|
||||
)
|
||||
if norm_type == "layer_norm":
|
||||
self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
|
||||
elif norm_type == "rms_norm":
|
||||
self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
|
||||
else:
|
||||
raise ValueError(f"unknown norm_type {norm_type}")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
conditioning_embedding: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
|
||||
if cu_seqlens is None:
|
||||
scale, shift = torch.chunk(emb, 2, dim=-1)
|
||||
x = self.norm(x) * (1 + scale) + shift
|
||||
else:
|
||||
sample_lens = cu_seqlens[1:] - cu_seqlens[:-1]
|
||||
flattened_x = x.view(-1, x.shape[-1])
|
||||
scale, shift = torch.chunk(emb, 2, dim=-1)
|
||||
scale_t = torch.repeat_interleave(scale, sample_lens, dim=0)
|
||||
shift_t = torch.repeat_interleave(shift, sample_lens, dim=0)
|
||||
flattened_x = self.norm(flattened_x) * (1 + scale_t) + shift_t
|
||||
x = flattened_x.view(x.shape)
|
||||
return x
|
||||
|
||||
|
||||
@dataclass
|
||||
class MageFlowParams:
|
||||
in_channels: int
|
||||
out_channels: int
|
||||
context_in_dim: int
|
||||
hidden_size: int
|
||||
num_heads: int
|
||||
depth: int
|
||||
axes_dim: list
|
||||
checkpoint: bool
|
||||
patch_size: int = 1
|
||||
|
||||
|
||||
class MageFlow(nn.Module):
|
||||
def __init__(self, params: MageFlowParams):
|
||||
super().__init__()
|
||||
self.params = params
|
||||
self.gradient_checkpointing = bool(params.checkpoint)
|
||||
self.in_channels = params.in_channels
|
||||
self.out_channels = params.out_channels
|
||||
self.inner_dim = params.hidden_size # num_attention_heads * attention_head_dim
|
||||
self.axes_dim = params.axes_dim
|
||||
self.num_attention_heads = params.num_heads
|
||||
self.attention_head_dim = self.inner_dim // self.num_attention_heads
|
||||
self.patch_size = params.patch_size
|
||||
assert sum(self.axes_dim) == self.attention_head_dim
|
||||
|
||||
self.pos_embed = MageFlowEmbedRope(
|
||||
theta=10000, axes_dim=self.axes_dim, scale_rope=True
|
||||
)
|
||||
self.img_in = nn.Linear(self.in_channels, self.inner_dim)
|
||||
self.txt_norm = RMSNorm(params.context_in_dim, eps=1e-6)
|
||||
self.txt_in = nn.Linear(params.context_in_dim, self.inner_dim)
|
||||
|
||||
self.time_text_embed = MageFlowTimestepProjEmbeddings(
|
||||
embedding_dim=self.inner_dim
|
||||
)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
MageFlowTransformerBlock(
|
||||
dim=self.inner_dim,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
attention_head_dim=self.attention_head_dim,
|
||||
)
|
||||
for _ in range(params.depth)
|
||||
]
|
||||
)
|
||||
|
||||
self.norm_out = AdaLayerNormContinuous(
|
||||
self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
|
||||
)
|
||||
self.proj_out = nn.Linear(
|
||||
self.inner_dim,
|
||||
self.patch_size * self.patch_size * self.out_channels,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
def enable_gradient_checkpointing(self, enable: bool = True):
|
||||
self.gradient_checkpointing = enable
|
||||
|
||||
def disable_gradient_checkpointing(self):
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def reset_rope(self):
|
||||
"""Rebuild the RoPE frequency tables. Needed after a meta-device init
|
||||
(the tables are plain tensors, not buffers, so ``load_state_dict`` /
|
||||
``assign=True`` never materializes them)."""
|
||||
self.pos_embed = MageFlowEmbedRope(
|
||||
theta=10000, axes_dim=self.axes_dim, scale_rope=True
|
||||
)
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return next(self.parameters()).device
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
txt: Tensor,
|
||||
timesteps: Tensor,
|
||||
img_shapes=None,
|
||||
img_cu_seqlens: Tensor = None,
|
||||
txt_cu_seqlens: Tensor = None,
|
||||
attention_kwargs: dict = None,
|
||||
) -> Tensor:
|
||||
if img.ndim != 3 or txt.ndim != 3:
|
||||
raise ValueError("Input img and txt tensors must have 3 dimensions.")
|
||||
|
||||
# Prepare vision RoPE (msrope); text tokens are not rotated.
|
||||
ms_pe = self.pos_embed(img_shapes, device=img.device)
|
||||
|
||||
img = self.img_in(img)
|
||||
txt = self.txt_norm(txt)
|
||||
|
||||
timesteps = timesteps.to(img.dtype)
|
||||
temb = self.time_text_embed(timesteps, img)
|
||||
|
||||
txt = self.txt_in(txt)
|
||||
# The reference adds a zero "text vector" to temb (vec_type is null in
|
||||
# the released config); kept for parity with the reference forward.
|
||||
txt_vec = torch.zeros(
|
||||
txt.shape[0], self.inner_dim, dtype=txt.dtype, device=txt.device
|
||||
)
|
||||
temb = temb + txt_vec
|
||||
|
||||
attention_kwargs = attention_kwargs or {}
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
if self.gradient_checkpointing and torch.is_grad_enabled():
|
||||
txt, img = checkpoint(
|
||||
block,
|
||||
img, # hidden_states
|
||||
txt, # encoder_hidden_states
|
||||
temb, # temb
|
||||
ms_pe, # image_rotary_emb
|
||||
txt_cu_seqlens, # txt_cu_lens
|
||||
img_cu_seqlens, # img_cu_lens
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
txt, img = block(
|
||||
hidden_states=img,
|
||||
encoder_hidden_states=txt,
|
||||
txt_cu_lens=txt_cu_seqlens,
|
||||
img_cu_lens=img_cu_seqlens,
|
||||
temb=temb,
|
||||
image_rotary_emb=ms_pe,
|
||||
joint_attention_kwargs=attention_kwargs,
|
||||
)
|
||||
|
||||
# Use only the image part (hidden_states) from the dual-stream blocks
|
||||
img = self.norm_out(img, temb, cu_seqlens=img_cu_seqlens)
|
||||
img = self.proj_out(img)
|
||||
return img
|
||||
703
extensions_built_in/diffusion_models/mageflow/src/vae.py
Normal file
703
extensions_built_in/diffusion_models/mageflow/src/vae.py
Normal file
@@ -0,0 +1,703 @@
|
||||
"""
|
||||
MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) wrapper.
|
||||
|
||||
Vendored from the reference implementation
|
||||
(github.com/microsoft/Mage, mage_flow/models/modules/mage_vae.py); loguru
|
||||
swapped for prints, otherwise unchanged. Loads the kl0.1 CoD ckpt layout used
|
||||
by ``vae/diffusion_pytorch_model.safetensors`` in the Mage-Flow HF repos:
|
||||
encoder weights: 'student.dconv_encoder.*' (packed mean+logvar, out_ch_mult=2)
|
||||
decoder weights: 'pipeline.*' (denoiser + y_embedder.decoder)
|
||||
|
||||
Latent shape: [B, 128, H/16, W/16] — no patch packing, no normalization.
|
||||
Encode input / decode output are [-1, 1] pixel images with H, W multiples of 16.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Primitive layers (vendored from GenCodec, inference subset)
|
||||
# ---------------------------------------------------------------------------
|
||||
def nonlinearity(x):
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
def Normalize(in_channels):
|
||||
return torch.nn.GroupNorm(
|
||||
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
|
||||
|
||||
def modulate(x, shift, scale):
|
||||
if x.dim() == 4:
|
||||
b, c = x.shape[:2]
|
||||
return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1)
|
||||
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
|
||||
class LayerNorm2d(nn.LayerNorm):
|
||||
def __init__(self, num_channels, eps=1e-6, affine=True):
|
||||
super().__init__(num_channels, eps=eps, elementwise_affine=affine)
|
||||
|
||||
def forward(self, x):
|
||||
# .contiguous() prevents a channels_last-strided NCHW view from
|
||||
# propagating into downstream depthwise convs, which would otherwise
|
||||
# hit a slow cuDNN path with a per-shape heuristic search.
|
||||
x = x.permute(0, 2, 3, 1).contiguous()
|
||||
x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
||||
return x.permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
|
||||
class _EncoderLayerNorm2d(LayerNorm2d):
|
||||
pass
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, x):
|
||||
in_dtype = x.dtype
|
||||
x = x.to(torch.float32)
|
||||
var = x.pow(2).mean(-1, keepdim=True)
|
||||
x = x * torch.rsqrt(var + self.variance_epsilon)
|
||||
return self.weight * x.to(in_dtype)
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""DConv-style timestep MLP (max_period=10000, freq_size=256, hidden=384)."""
|
||||
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
|
||||
).to(t.device)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
|
||||
return emb
|
||||
|
||||
def forward(self, t):
|
||||
emb = self.timestep_embedding(t, self.frequency_embedding_size)
|
||||
return self.mlp(emb.to(self.mlp[0].weight.dtype))
|
||||
|
||||
|
||||
class BottleneckPatchEmbed(nn.Module):
|
||||
"""Image patch embed concatenated with a per-patch conditioning vector."""
|
||||
|
||||
def __init__(
|
||||
self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True
|
||||
):
|
||||
super().__init__()
|
||||
self.proj1 = nn.Conv2d(
|
||||
in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False
|
||||
)
|
||||
self.proj2 = nn.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
|
||||
|
||||
def forward(self, x, cond):
|
||||
return self.proj2(torch.cat([self.proj1(x), cond], dim=1))
|
||||
|
||||
|
||||
class DiCoBlock(nn.Module):
|
||||
"""DConv block with adaLN modulation."""
|
||||
|
||||
def __init__(self, hidden_size, mlp_ratio=4.0):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
|
||||
self.conv2 = nn.Conv2d(
|
||||
hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True
|
||||
)
|
||||
self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
|
||||
|
||||
self.ca = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d(1),
|
||||
nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
ffn = int(mlp_ratio * hidden_size)
|
||||
self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
|
||||
self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
|
||||
|
||||
self.norm1 = LayerNorm2d(hidden_size, affine=False)
|
||||
self.norm2 = LayerNorm2d(hidden_size, affine=False)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, 6 * hidden_size, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, inp, c):
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.adaLN_modulation(c).chunk(6, dim=1)
|
||||
)
|
||||
x = modulate(self.norm1(inp), shift_msa, scale_msa)
|
||||
x = F.gelu(self.conv2(self.conv1(x)))
|
||||
x = x * self.ca(x)
|
||||
x = self.conv3(x)
|
||||
x = inp + gate_msa[..., None, None] * x
|
||||
x = x + gate_mlp[..., None, None] * self.conv5(
|
||||
F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp)))
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class _EncoderDiCoBlock(nn.Module):
|
||||
"""DiCoBlock without adaLN, for the encoder pathway."""
|
||||
|
||||
def __init__(self, hidden_size, mlp_ratio=4.0):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
|
||||
self.conv2 = nn.Conv2d(
|
||||
hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True
|
||||
)
|
||||
self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
|
||||
self.ca = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d(1),
|
||||
nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
ffn = int(mlp_ratio * hidden_size)
|
||||
self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
|
||||
self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
|
||||
self.norm1 = _EncoderLayerNorm2d(hidden_size)
|
||||
self.norm2 = _EncoderLayerNorm2d(hidden_size)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.norm1(inp)
|
||||
x = F.gelu(self.conv2(self.conv1(x)))
|
||||
x = x * self.ca(x)
|
||||
x = self.conv3(x)
|
||||
x = inp + x
|
||||
return x + self.conv5(F.gelu(self.conv4(self.norm2(x))))
|
||||
|
||||
|
||||
class NerfEmbedder(nn.Module):
|
||||
"""Patch-position embedder used by the DConv decoder x-pathway."""
|
||||
|
||||
def __init__(self, in_channels, hidden_size_input, max_freqs=8):
|
||||
super().__init__()
|
||||
self.max_freqs = max_freqs
|
||||
self.embedder = nn.Sequential(
|
||||
nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True),
|
||||
)
|
||||
|
||||
@lru_cache
|
||||
def fetch_pos(self, patch_size, device, dtype):
|
||||
pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
|
||||
pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
|
||||
pos_x = pos_x.reshape(-1, 1, 1)
|
||||
pos_y = pos_y.reshape(-1, 1, 1)
|
||||
freqs = torch.linspace(
|
||||
0, self.max_freqs, self.max_freqs, dtype=dtype, device=device
|
||||
)
|
||||
fx = freqs[None, :, None]
|
||||
fy = freqs[None, None, :]
|
||||
coeffs = (1 + fx * fy) ** -1
|
||||
dct_x = torch.cos(pos_x * fx * torch.pi)
|
||||
dct_y = torch.cos(pos_y * fy * torch.pi)
|
||||
return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2)
|
||||
|
||||
def forward(self, x):
|
||||
B, P2, _ = x.shape
|
||||
ps = int(P2**0.5)
|
||||
dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1)
|
||||
return self.embedder(torch.cat([x, dct], dim=-1))
|
||||
|
||||
|
||||
class NerfFinalLayer(nn.Module):
|
||||
def __init__(self, hidden_size, out_channels):
|
||||
super().__init__()
|
||||
self.norm = RMSNorm(hidden_size)
|
||||
self.linear = nn.Linear(hidden_size, out_channels, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.linear(self.norm(x))
|
||||
|
||||
|
||||
class SimpleMLPAdaLN(nn.Module):
|
||||
"""Final small MLP that maps NerfEmbedder features to per-patch RGB."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
model_channels,
|
||||
out_channels,
|
||||
z_channels,
|
||||
num_res_blocks,
|
||||
patch_size,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.patch_size = patch_size
|
||||
|
||||
self.cond_embed = nn.Linear(z_channels, patch_size**2 * model_channels)
|
||||
self.input_proj = nn.Linear(in_channels, model_channels)
|
||||
|
||||
self.res_blocks = nn.ModuleList(
|
||||
_MLPResBlock(model_channels) for _ in range(num_res_blocks)
|
||||
)
|
||||
|
||||
def forward(self, x, c):
|
||||
x = self.input_proj(x)
|
||||
c = self.cond_embed(c).reshape(c.shape[0], self.patch_size**2, -1)
|
||||
for block in self.res_blocks:
|
||||
x = block(x, c)
|
||||
return x
|
||||
|
||||
|
||||
class _MLPResBlock(nn.Module):
|
||||
def __init__(self, channels):
|
||||
super().__init__()
|
||||
self.in_ln = nn.LayerNorm(channels, eps=1e-6)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(channels, channels, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(channels, channels, bias=True),
|
||||
)
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(channels, 3 * channels, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, x, y):
|
||||
shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
|
||||
h = self.in_ln(x) * (1 + scale) + shift
|
||||
return x + gate * self.mlp(h)
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
"""GroupNorm + Conv ResBlock used by the CoD Decoder."""
|
||||
|
||||
def __init__(self, *, in_channels, out_channels=None, dropout=0.0):
|
||||
super().__init__()
|
||||
out_channels = out_channels or in_channels
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.norm1 = Normalize(in_channels)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
|
||||
self.norm2 = Normalize(out_channels)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
|
||||
if in_channels != out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.conv1(nonlinearity(self.norm1(x)))
|
||||
h = self.conv2(self.dropout(nonlinearity(self.norm2(h))))
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
return x + h
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
"""Patched self-attention used at inference (eval mode of the original)."""
|
||||
|
||||
def __init__(self, in_channels, patch_size=32):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.patch_size = patch_size
|
||||
self.norm = Normalize(in_channels)
|
||||
self.q = nn.Conv2d(in_channels, in_channels, 1)
|
||||
self.k = nn.Conv2d(in_channels, in_channels, 1)
|
||||
self.v = nn.Conv2d(in_channels, in_channels, 1)
|
||||
self.proj_out = nn.Conv2d(in_channels, in_channels, 1)
|
||||
|
||||
def forward(self, x):
|
||||
h_ = self.norm(x)
|
||||
Q = self.q(h_)
|
||||
K = self.k(h_)
|
||||
V = self.v(h_)
|
||||
|
||||
d = self.patch_size
|
||||
b, c, H, W = Q.shape
|
||||
pad_h = (d - H % d) % d
|
||||
pad_w = (d - W % d) % d
|
||||
if pad_h or pad_w:
|
||||
Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate")
|
||||
K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate")
|
||||
V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate")
|
||||
_, _, H_pad, W_pad = Q.shape
|
||||
nph, npw = H_pad // d, W_pad // d
|
||||
np_ = nph * npw
|
||||
|
||||
def to_patches(t):
|
||||
return (
|
||||
t.reshape(b, c, nph, d, npw, d)
|
||||
.permute(0, 2, 4, 1, 3, 5)
|
||||
.reshape(b * np_, c, d * d)
|
||||
)
|
||||
|
||||
Q = to_patches(Q)
|
||||
K = to_patches(K)
|
||||
V = to_patches(V)
|
||||
|
||||
w_ = torch.bmm(Q.permute(0, 2, 1), K) * (c**-0.5)
|
||||
w_ = F.softmax(w_, dim=2).permute(0, 2, 1)
|
||||
h_ = (
|
||||
torch.bmm(V, w_)
|
||||
.reshape(b, nph, npw, c, d, d)
|
||||
.permute(0, 3, 1, 4, 2, 5)
|
||||
.reshape(b, c, H_pad, W_pad)
|
||||
)
|
||||
if pad_h or pad_w:
|
||||
h_ = h_[:, :, :H, :W]
|
||||
return x + self.proj_out(h_)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# adaLN constant-folding: at fixed t=0, adaLN_modulation(c) is constant.
|
||||
# Replace the MLP with a buffer so DiCoBlock.forward stays unchanged.
|
||||
# ---------------------------------------------------------------------------
|
||||
class _ConstAdaLN(nn.Module):
|
||||
def __init__(self, modulation: torch.Tensor):
|
||||
super().__init__()
|
||||
self.register_buffer("modulation", modulation.detach().clone())
|
||||
|
||||
def forward(self, c):
|
||||
b = c.shape[0]
|
||||
if self.modulation.shape[0] != b:
|
||||
return self.modulation.expand(b, *self.modulation.shape[1:])
|
||||
return self.modulation
|
||||
|
||||
|
||||
def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int:
|
||||
# Only DiCoBlock is targeted: its adaLN is conditioned solely on t.
|
||||
# Other adaLN_modulation submodules (e.g. _MLPResBlock in the decoder MLP)
|
||||
# take a per-position latent and must not be folded.
|
||||
n = 0
|
||||
for child in module.modules():
|
||||
if not isinstance(child, DiCoBlock):
|
||||
continue
|
||||
adaln = child.adaLN_modulation
|
||||
if isinstance(adaln, _ConstAdaLN):
|
||||
continue
|
||||
with torch.no_grad():
|
||||
mod = adaln(c)
|
||||
child.adaLN_modulation = _ConstAdaLN(mod)
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CoD Decoder: latent → conditioning features for the denoiser
|
||||
# ---------------------------------------------------------------------------
|
||||
class _Decoder(nn.Module):
|
||||
"""ds=16, up2x=True, light=True only."""
|
||||
|
||||
def __init__(self, out_ch=384, z_ch=128):
|
||||
super().__init__()
|
||||
self.conv_in = nn.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
self.block = nn.Sequential(
|
||||
ResnetBlock(in_channels=out_ch, out_channels=out_ch),
|
||||
AttnBlock(out_ch, patch_size=32),
|
||||
ResnetBlock(in_channels=out_ch, out_channels=out_ch),
|
||||
AttnBlock(out_ch, patch_size=32),
|
||||
ResnetBlock(in_channels=out_ch, out_channels=out_ch),
|
||||
)
|
||||
self.norm_out = Normalize(out_ch)
|
||||
self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
self.ada = nn.Identity()
|
||||
|
||||
def forward(self, z):
|
||||
h = self.block(self.conv_in(z))
|
||||
h = self.conv_out(nonlinearity(self.norm_out(h)))
|
||||
return self.ada(h)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DConvEncoder: image → packed (mean, logvar) latent
|
||||
# ---------------------------------------------------------------------------
|
||||
class _DConvEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
z_ch=128,
|
||||
hidden_size=384,
|
||||
num_blocks=21,
|
||||
patch_size=16,
|
||||
mlp_ratio=4.0,
|
||||
head_size=768,
|
||||
num_head_blocks=2,
|
||||
out_ch_mult=2,
|
||||
):
|
||||
super().__init__()
|
||||
self.z_ch = z_ch
|
||||
self.patch_size = patch_size
|
||||
self.patch_cond_embed = nn.Conv2d(
|
||||
3, head_size, kernel_size=patch_size, stride=patch_size, bias=True
|
||||
)
|
||||
self.head_blocks = nn.ModuleList(
|
||||
[
|
||||
_EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio)
|
||||
for _ in range(num_head_blocks)
|
||||
]
|
||||
)
|
||||
self.proj_down = nn.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
|
||||
self.z_proj = nn.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True)
|
||||
self.fuse_proj = nn.Conv2d(
|
||||
hidden_size * 2, hidden_size, kernel_size=1, bias=True
|
||||
)
|
||||
self.t_embedder = TimestepEmbedder(hidden_size)
|
||||
self.blocks = nn.ModuleList(
|
||||
[DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)]
|
||||
)
|
||||
self.norm_out = LayerNorm2d(hidden_size)
|
||||
self.proj_out = nn.Conv2d(
|
||||
hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True
|
||||
)
|
||||
|
||||
def forward_pred(self, z_t, t, y):
|
||||
cond = self.patch_cond_embed(y)
|
||||
for block in self.head_blocks:
|
||||
cond = block(cond)
|
||||
cond = self.proj_down(cond)
|
||||
|
||||
s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1))
|
||||
c = self.t_embedder(t.view(-1))
|
||||
for block in self.blocks:
|
||||
s = block(s, c)
|
||||
return self.proj_out(self.norm_out(s))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DConv denoiser: latent (via cond) + zero noise → reconstructed image
|
||||
# ---------------------------------------------------------------------------
|
||||
class _YEmbedder(nn.Module):
|
||||
"""Holds only the CoD decoder; the original Flux2 VAE encoder side is omitted."""
|
||||
|
||||
def __init__(self, ch=384, z_ch=128):
|
||||
super().__init__()
|
||||
self.decoder = _Decoder(out_ch=ch, z_ch=z_ch)
|
||||
|
||||
|
||||
class _DConvDenoiser(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
patch_size=16,
|
||||
in_channels=3,
|
||||
hidden_size=384,
|
||||
hidden_size_x=32,
|
||||
mlp_ratio=4.0,
|
||||
num_blocks=24,
|
||||
num_cond_blocks=21,
|
||||
bottleneck_dim=128,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.patch_size = patch_size
|
||||
self.hidden_size = hidden_size
|
||||
self.num_cond_blocks = num_cond_blocks
|
||||
|
||||
self.t_embedder = TimestepEmbedder(hidden_size)
|
||||
self.y_embedder_x = nn.Conv2d(
|
||||
hidden_size, hidden_size_x * patch_size**2, 1, 1, 0
|
||||
)
|
||||
self.x_embedder = NerfEmbedder(
|
||||
in_channels + hidden_size_x, hidden_size_x, max_freqs=8
|
||||
)
|
||||
self.s_embedder = BottleneckPatchEmbed(
|
||||
patch_size, in_channels, bottleneck_dim, hidden_size, bias=True
|
||||
)
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
DiCoBlock(hidden_size, mlp_ratio=mlp_ratio)
|
||||
for _ in range(num_cond_blocks)
|
||||
]
|
||||
)
|
||||
self.dec_net = SimpleMLPAdaLN(
|
||||
in_channels=hidden_size_x,
|
||||
model_channels=hidden_size_x,
|
||||
out_channels=in_channels,
|
||||
z_channels=hidden_size,
|
||||
num_res_blocks=num_blocks - num_cond_blocks,
|
||||
patch_size=patch_size,
|
||||
)
|
||||
self.final_layer = NerfFinalLayer(hidden_size_x, in_channels)
|
||||
self.y_embedder = _YEmbedder(ch=hidden_size, z_ch=bottleneck_dim)
|
||||
|
||||
def forward(self, x, t, cond):
|
||||
b, _, h, w = x.shape
|
||||
c = self.t_embedder(t.view(-1))
|
||||
|
||||
s = self.s_embedder(x, cond)
|
||||
for block in self.blocks:
|
||||
s = block(s, c)
|
||||
|
||||
length = s.shape[-2] * s.shape[-1]
|
||||
s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
|
||||
|
||||
x = torch.nn.functional.unfold(
|
||||
x, kernel_size=self.patch_size, stride=self.patch_size
|
||||
)
|
||||
x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1)
|
||||
x = (
|
||||
x.reshape(b, -1, self.patch_size**2, length)
|
||||
.permute(0, 3, 2, 1)
|
||||
.flatten(0, 1)
|
||||
)
|
||||
x = self.x_embedder(x)
|
||||
|
||||
x = self.dec_net(x, s)
|
||||
x = self.final_layer(x)
|
||||
x = x.transpose(1, 2).reshape(b, length, -1)
|
||||
return torch.nn.functional.fold(
|
||||
x.transpose(1, 2).contiguous(),
|
||||
(h, w),
|
||||
kernel_size=self.patch_size,
|
||||
stride=self.patch_size,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Wrapper
|
||||
# ---------------------------------------------------------------------------
|
||||
def _load_state_dict(ckpt_path: str):
|
||||
if ckpt_path.endswith(".safetensors"):
|
||||
from safetensors.torch import load_file
|
||||
|
||||
return load_file(ckpt_path, device="cpu")
|
||||
if os.path.exists(os.path.join(ckpt_path, "checkpoint-state_dict.pt")):
|
||||
ckpt_path = os.path.join(ckpt_path, "checkpoint-state_dict.pt")
|
||||
elif os.path.isdir(ckpt_path):
|
||||
ckpt_path = os.path.join(ckpt_path, "checkpoint", "mp_rank_00_model_states.pt")
|
||||
state = torch.load(ckpt_path, map_location="cpu")
|
||||
if "module" in state:
|
||||
return state["module"]
|
||||
if "state_dict" in state:
|
||||
return state["state_dict"]
|
||||
return state
|
||||
|
||||
|
||||
class MageVAE(nn.Module):
|
||||
"""
|
||||
Encode: DConvEncoder (one-step diffusion) → latent [B, 128, H/16, W/16]
|
||||
Decode: DConvDenoiser + CoD Decoder → image [B, 3, H, W] in [-1, 1]
|
||||
"""
|
||||
|
||||
latent_channels = 128
|
||||
downsample_factor = 16
|
||||
|
||||
def __init__(self, ckpt_path: str, sample_posterior: bool = True):
|
||||
super().__init__()
|
||||
self.sample_posterior = sample_posterior
|
||||
|
||||
self.dconv_encoder = _DConvEncoder()
|
||||
self.decoder_model = _DConvDenoiser()
|
||||
|
||||
sd = _load_state_dict(ckpt_path)
|
||||
self._load_encoder(sd, ckpt_path)
|
||||
self._load_decoder(sd, ckpt_path)
|
||||
|
||||
# adaLN modulation depends only on t, and we always run at t=0.
|
||||
# Precompute and drop the MLPs once at construction (~37M params saved).
|
||||
self._freeze_adaln_cache()
|
||||
|
||||
def _load_encoder(self, sd, ckpt_path):
|
||||
prefix = "student.dconv_encoder."
|
||||
enc_sd = {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)}
|
||||
if not enc_sd:
|
||||
raise RuntimeError(f"MageVAE encoder: no '{prefix}*' keys in {ckpt_path}")
|
||||
proj = enc_sd.get("proj_out.weight")
|
||||
if proj is None or proj.shape[0] != 2 * self.latent_channels:
|
||||
raise RuntimeError(
|
||||
f"MageVAE encoder: expected packed mean+logvar (proj_out out_channels="
|
||||
f"{2 * self.latent_channels}), got {None if proj is None else tuple(proj.shape)}"
|
||||
)
|
||||
missing, unexpected = self.dconv_encoder.load_state_dict(enc_sd, strict=False)
|
||||
if missing:
|
||||
print(f"MageVAE encoder missing keys: {missing[:10]}")
|
||||
|
||||
def _load_decoder(self, sd, ckpt_path):
|
||||
prefix = "pipeline."
|
||||
if not any(k.startswith(prefix) for k in sd):
|
||||
raise RuntimeError(f"MageVAE decoder: no '{prefix}*' keys in {ckpt_path}")
|
||||
model_dict = self.decoder_model.state_dict()
|
||||
matched = {}
|
||||
for k, v in sd.items():
|
||||
if not k.startswith(prefix):
|
||||
continue
|
||||
new_k = k[len(prefix) :]
|
||||
if new_k.startswith("y_embedder.encoder.") or new_k.startswith(
|
||||
"y_embedder.bottleneck."
|
||||
):
|
||||
continue
|
||||
if new_k in model_dict and model_dict[new_k].shape == v.shape:
|
||||
matched[new_k] = v
|
||||
self.decoder_model.load_state_dict(matched, strict=False)
|
||||
if not matched:
|
||||
raise RuntimeError(f"MageVAE decoder: 0 params matched from {ckpt_path}")
|
||||
|
||||
@torch.no_grad()
|
||||
def _moments(self, x: torch.Tensor):
|
||||
B, _, H, W = x.shape
|
||||
ps = self.dconv_encoder.patch_size
|
||||
z_t = torch.zeros(
|
||||
B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype
|
||||
)
|
||||
t = torch.zeros(B, device=x.device, dtype=x.dtype)
|
||||
out = self.dconv_encoder.forward_pred(z_t, t, x)
|
||||
mean = out[:, : self.latent_channels]
|
||||
logvar = out[:, self.latent_channels :].clamp(min=-20.0, max=10.0)
|
||||
return mean, logvar
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, x: torch.Tensor) -> torch.Tensor:
|
||||
ps = self.dconv_encoder.patch_size
|
||||
H, W = x.shape[-2], x.shape[-1]
|
||||
if H % ps or W % ps:
|
||||
raise ValueError(f"H, W must be multiples of {ps}, got ({H}, {W})")
|
||||
mean, logvar = self._moments(x)
|
||||
if self.sample_posterior:
|
||||
return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean)
|
||||
return mean
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, z: torch.Tensor) -> torch.Tensor:
|
||||
cond = self.decoder_model.y_embedder.decoder(z)
|
||||
B = z.shape[0]
|
||||
H = z.shape[2] * self.downsample_factor
|
||||
W = z.shape[3] * self.downsample_factor
|
||||
noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype)
|
||||
t = torch.zeros(B, device=z.device, dtype=z.dtype)
|
||||
return self.decoder_model.forward(noise, t, cond)
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return next(self.parameters()).device
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
def _freeze_adaln_cache(self):
|
||||
"""Constant-fold adaLN_modulation MLPs at t=0 (encoder + decoder)."""
|
||||
device = next(self.parameters()).device
|
||||
dtype = next(self.parameters()).dtype
|
||||
t = torch.zeros(1, device=device, dtype=dtype)
|
||||
c_enc = self.dconv_encoder.t_embedder(t)
|
||||
_replace_adaln_with_const(self.dconv_encoder, c_enc)
|
||||
c_dec = self.decoder_model.t_embedder(t)
|
||||
_replace_adaln_with_const(self.decoder_model, c_dec)
|
||||
Reference in New Issue
Block a user