From 2ca92088621f9f191bf86f0069a5ac712ed80e95 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Tue, 6 Aug 2024 15:24:46 +0800
Subject: [PATCH] Add files via upload

add
---
 k_diffusion/models/axial_rope.py           | 109 +++
 k_diffusion/models/flags.py                |  56 ++
 k_diffusion/models/flops.py                |  54 ++
 k_diffusion/models/image_transformer_v1.py | 344 ++++++++++
 k_diffusion/models/image_transformer_v2.py | 753 +++++++++++++++++++++
 k_diffusion/models/image_v1.py             | 176 +++++
 6 files changed, 1492 insertions(+)
 create mode 100644 k_diffusion/models/axial_rope.py
 create mode 100644 k_diffusion/models/flags.py
 create mode 100644 k_diffusion/models/flops.py
 create mode 100644 k_diffusion/models/image_transformer_v1.py
 create mode 100644 k_diffusion/models/image_transformer_v2.py
 create mode 100644 k_diffusion/models/image_v1.py

diff --git a/k_diffusion/models/axial_rope.py b/k_diffusion/models/axial_rope.py
new file mode 100644
index 0000000..cadc3b6
--- /dev/null
+++ b/k_diffusion/models/axial_rope.py
@@ -0,0 +1,109 @@
+import math
+
+import torch
+import torch._dynamo
+from torch import nn
+
+from . import flags
+
+if flags.get_use_compile():
+    torch._dynamo.config.suppress_errors = True
+
+
+def rotate_half(x):
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    x = torch.stack((-x2, x1), dim=-1)
+    *shape, d, r = x.shape
+    return x.view(*shape, d * r)
+
+
+@flags.compile_wrap
+def apply_rotary_emb(freqs, t, start_index=0, scale=1.0):
+    freqs = freqs.to(t)
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+    assert rot_dim <= t.shape[-1], f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
+    t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
+    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
+    return torch.cat((t_left, t, t_right), dim=-1)
+
+
+def centers(start, stop, num, dtype=None, device=None):
+    edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device)
+    return (edges[:-1] + edges[1:]) / 2
+
+
+def make_grid(h_pos, w_pos):
+    grid = torch.stack(torch.meshgrid(h_pos, w_pos, indexing='ij'), dim=-1)
+    h, w, d = grid.shape
+    return grid.view(h * w, d)
+
+
+def bounding_box(h, w, pixel_aspect_ratio=1.0):
+    # Adjusted dimensions
+    w_adj = w
+    h_adj = h * pixel_aspect_ratio
+
+    # Adjusted aspect ratio
+    ar_adj = w_adj / h_adj
+
+    # Determine bounding box based on the adjusted aspect ratio
+    y_min, y_max, x_min, x_max = -1.0, 1.0, -1.0, 1.0
+    if ar_adj > 1:
+        y_min, y_max = -1 / ar_adj, 1 / ar_adj
+    elif ar_adj < 1:
+        x_min, x_max = -ar_adj, ar_adj
+
+    return y_min, y_max, x_min, x_max
+
+
+def make_axial_pos(h, w, pixel_aspect_ratio=1.0, align_corners=False, dtype=None, device=None):
+    y_min, y_max, x_min, x_max = bounding_box(h, w, pixel_aspect_ratio)
+    if align_corners:
+        h_pos = torch.linspace(y_min, y_max, h, dtype=dtype, device=device)
+        w_pos = torch.linspace(x_min, x_max, w, dtype=dtype, device=device)
+    else:
+        h_pos = centers(y_min, y_max, h, dtype=dtype, device=device)
+        w_pos = centers(x_min, x_max, w, dtype=dtype, device=device)
+    return make_grid(h_pos, w_pos)
+
+
+def freqs_pixel(max_freq=10.0):
+    def init(shape):
+        freqs = torch.linspace(1.0, max_freq / 2, shape[-1]) * math.pi
+        return freqs.log().expand(shape)
+    return init
+
+
+def freqs_pixel_log(max_freq=10.0):
+    def init(shape):
+        log_min = math.log(math.pi)
+        log_max = math.log(max_freq * math.pi / 2)
+        return torch.linspace(log_min, log_max, shape[-1]).expand(shape)
+    return init
+
+
+class AxialRoPE(nn.Module):
+    def __init__(self, dim, n_heads, start_index=0, freqs_init=freqs_pixel_log(max_freq=10.0)):
+        super().__init__()
+        self.n_heads = n_heads
+        self.start_index = start_index
+        log_freqs = freqs_init((n_heads, dim // 4))
+        self.freqs_h = nn.Parameter(log_freqs.clone())
+        self.freqs_w = nn.Parameter(log_freqs.clone())
+
+    def extra_repr(self):
+        dim = (self.freqs_h.shape[-1] + self.freqs_w.shape[-1]) * 2
+        return f"dim={dim}, n_heads={self.n_heads}, start_index={self.start_index}"
+
+    def get_freqs(self, pos):
+        if pos.shape[-1] != 2:
+            raise ValueError("input shape must be (..., 2)")
+        freqs_h = pos[..., None, None, 0] * self.freqs_h.exp()
+        freqs_w = pos[..., None, None, 1] * self.freqs_w.exp()
+        freqs = torch.cat((freqs_h, freqs_w), dim=-1).repeat_interleave(2, dim=-1)
+        return freqs.transpose(-2, -3)
+
+    def forward(self, x, pos):
+        freqs = self.get_freqs(pos)
+        return apply_rotary_emb(freqs, x, self.start_index)
diff --git a/k_diffusion/models/flags.py b/k_diffusion/models/flags.py
new file mode 100644
index 0000000..6d555b7
--- /dev/null
+++ b/k_diffusion/models/flags.py
@@ -0,0 +1,56 @@
+from contextlib import contextmanager
+from functools import update_wrapper
+import os
+import threading
+
+import torch
+
+
+def get_use_compile():
+    return os.environ.get("K_DIFFUSION_USE_COMPILE", "1") == "1"
+
+
+def get_use_flash_attention_2():
+    return os.environ.get("K_DIFFUSION_USE_FLASH_2", "1") == "1"
+
+
+state = threading.local()
+state.checkpointing = False
+
+
+@contextmanager
+def checkpointing(enable=True):
+    try:
+        old_checkpointing, state.checkpointing = state.checkpointing, enable
+        yield
+    finally:
+        state.checkpointing = old_checkpointing
+
+
+def get_checkpointing():
+    return getattr(state, "checkpointing", False)
+
+
+class compile_wrap:
+    def __init__(self, function, *args, **kwargs):
+        self.function = function
+        self.args = args
+        self.kwargs = kwargs
+        self._compiled_function = None
+        update_wrapper(self, function)
+
+    @property
+    def compiled_function(self):
+        if self._compiled_function is not None:
+            return self._compiled_function
+        if get_use_compile():
+            try:
+                self._compiled_function = torch.compile(self.function, *self.args, **self.kwargs)
+            except RuntimeError:
+                self._compiled_function = self.function
+        else:
+            self._compiled_function = self.function
+        return self._compiled_function
+
+    def __call__(self, *args, **kwargs):
+        return self.compiled_function(*args, **kwargs)
diff --git a/k_diffusion/models/flops.py b/k_diffusion/models/flops.py
new file mode 100644
index 0000000..e98e2a8
--- /dev/null
+++ b/k_diffusion/models/flops.py
@@ -0,0 +1,54 @@
+from contextlib import contextmanager
+import math
+import threading
+
+
+state = threading.local()
+state.flop_counter = None
+
+
+@contextmanager
+def flop_counter(enable=True):
+    try:
+        old_flop_counter = state.flop_counter
+        state.flop_counter = FlopCounter() if enable else None
+        yield state.flop_counter
+    finally:
+        state.flop_counter = old_flop_counter
+
+
+class FlopCounter:
+    def __init__(self):
+        self.ops = []
+
+    def op(self, op, *args, **kwargs):
+        self.ops.append((op, args, kwargs))
+
+    @property
+    def flops(self):
+        flops = 0
+        for op, args, kwargs in self.ops:
+            flops += op(*args, **kwargs)
+        return flops
+
+
+def op(op, *args, **kwargs):
+    if getattr(state, "flop_counter", None):
+        state.flop_counter.op(op, *args, **kwargs)
+
+
+def op_linear(x, weight):
+    return math.prod(x) * weight[0]
+
+
+def op_attention(q, k, v):
+    *b, s_q, d_q = q
+    *b, s_k, d_k = k
+    *b, s_v, d_v = v
+    return math.prod(b) * s_q * s_k * (d_q + d_v)
+
+
+def op_natten(q, k, v, kernel_size):
+    *q_rest, d_q = q
+    *_, d_v = v
+    return math.prod(q_rest) * (d_q + d_v) * kernel_size**2
diff --git a/k_diffusion/models/image_transformer_v1.py b/k_diffusion/models/image_transformer_v1.py
new file mode 100644
index 0000000..c37c639
--- /dev/null
+++ b/k_diffusion/models/image_transformer_v1.py
@@ -0,0 +1,344 @@
+"""k-diffusion transformer diffusion models, version 1."""
+
+import math
+
+from einops import rearrange
+import torch
+from torch import nn
+import torch._dynamo
+from torch.nn import functional as F
+
+from . import flags
+from .. import layers
+from .axial_rope import AxialRoPE, make_axial_pos
+
+if flags.get_use_compile():
+    torch._dynamo.config.suppress_errors = True
+
+
+def zero_init(layer):
+    nn.init.zeros_(layer.weight)
+    if layer.bias is not None:
+        nn.init.zeros_(layer.bias)
+    return layer
+
+
+def checkpoint_helper(function, *args, **kwargs):
+    if flags.get_checkpointing():
+        kwargs.setdefault("use_reentrant", True)
+        return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+    else:
+        return function(*args, **kwargs)
+
+
+def tag_param(param, tag):
+    if not hasattr(param, "_tags"):
+        param._tags = set([tag])
+    else:
+        param._tags.add(tag)
+    return param
+
+
+def tag_module(module, tag):
+    for param in module.parameters():
+        tag_param(param, tag)
+    return module
+
+
+def apply_wd(module):
+    for name, param in module.named_parameters():
+        if name.endswith("weight"):
+            tag_param(param, "wd")
+    return module
+
+
+def filter_params(function, module):
+    for param in module.parameters():
+        tags = getattr(param, "_tags", set())
+        if function(tags):
+            yield param
+
+
+def scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0):
+    if flags.get_use_flash_attention_2() and attn_mask is None:
+        try:
+            from flash_attn import flash_attn_func
+            q_ = q.transpose(-3, -2)
+            k_ = k.transpose(-3, -2)
+            v_ = v.transpose(-3, -2)
+            o_ = flash_attn_func(q_, k_, v_, dropout_p=dropout_p)
+            return o_.transpose(-3, -2)
+        except (ImportError, RuntimeError):
+            pass
+    return F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p=dropout_p)
+
+
+@flags.compile_wrap
+def geglu(x):
+    a, b = x.chunk(2, dim=-1)
+    return a * F.gelu(b)
+
+
+@flags.compile_wrap
+def rms_norm(x, scale, eps):
+    dtype = torch.promote_types(x.dtype, torch.float32)
+    mean_sq = torch.mean(x.to(dtype)**2, dim=-1, keepdim=True)
+    scale = scale.to(dtype) * torch.rsqrt(mean_sq + eps)
+    return x * scale.to(x.dtype)
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        return geglu(x)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, param_shape, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(param_shape))
+
+    def extra_repr(self):
+        return f"shape={tuple(self.scale.shape)}, eps={self.eps}"
+
+    def forward(self, x):
+        return rms_norm(x, self.scale, self.eps)
+
+
+class QKNorm(nn.Module):
+    def __init__(self, n_heads, eps=1e-6, max_scale=100.0):
+        super().__init__()
+        self.eps = eps
+        self.max_scale = math.log(max_scale)
+        self.scale = nn.Parameter(torch.full((n_heads,), math.log(10.0)))
+        self.proj_()
+
+    def extra_repr(self):
+        return f"n_heads={self.scale.shape[0]}, eps={self.eps}"
+
+    @torch.no_grad()
+    def proj_(self):
+        """Modify the scale in-place so it doesn't get "stuck" with zero gradient if it's clamped
+        to the max value."""
+        self.scale.clamp_(max=self.max_scale)
+
+    def forward(self, x):
+        self.proj_()
+        scale = torch.exp(0.5 * self.scale - 0.25 * math.log(x.shape[-1]))
+        return rms_norm(x, scale[:, None, None], self.eps)
+
+
+class AdaRMSNorm(nn.Module):
+    def __init__(self, features, cond_features, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.linear = apply_wd(zero_init(nn.Linear(cond_features, features, bias=False)))
+        tag_module(self.linear, "mapping")
+
+    def extra_repr(self):
+        return f"eps={self.eps},"
+
+    def forward(self, x, cond):
+        return rms_norm(x, self.linear(cond) + 1, self.eps)
+
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.norm = AdaRMSNorm(d_model, d_model)
+        self.qkv_proj = apply_wd(nn.Linear(d_model, d_model * 3, bias=False))
+        self.qk_norm = QKNorm(self.n_heads)
+        self.pos_emb = AxialRoPE(d_head, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(nn.Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head},"
+
+    def forward(self, x, pos, attn_mask, cond):
+        skip = x
+        x = self.norm(x, cond)
+        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
+        q = rearrange(q, "n l (h e) -> n h l e", e=self.d_head)
+        k = rearrange(k, "n l (h e) -> n h l e", e=self.d_head)
+        v = rearrange(v, "n l (h e) -> n h l e", e=self.d_head)
+        q = self.pos_emb(self.qk_norm(q), pos)
+        k = self.pos_emb(self.qk_norm(k), pos)
+        x = scaled_dot_product_attention(q, k, v, attn_mask)
+        x = rearrange(x, "n h l e -> n l (h e)")
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class FeedForwardBlock(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.norm = AdaRMSNorm(d_model, d_model)
+        self.up_proj = apply_wd(nn.Linear(d_model, d_ff * 2, bias=False))
+        self.act = GEGLU()
+        self.dropout = nn.Dropout(dropout)
+        self.down_proj = apply_wd(zero_init(nn.Linear(d_ff, d_model, bias=False)))
+
+    def forward(self, x, cond):
+        skip = x
+        x = self.norm(x, cond)
+        x = self.up_proj(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.down_proj(x)
+        return x + skip
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, dropout=0.0):
+        super().__init__()
+        self.self_attn = SelfAttentionBlock(d_model, d_head, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, dropout=dropout)
+
+    def forward(self, x, pos, attn_mask, cond):
+        x = checkpoint_helper(self.self_attn, x, pos, attn_mask, cond)
+        x = checkpoint_helper(self.ff, x, cond)
+        return x
+
+
+class Patching(nn.Module):
+    def __init__(self, features, patch_size):
+        super().__init__()
+        self.features = features
+        self.patch_size = patch_size
+        self.d_out = features * patch_size[0] * patch_size[1]
+
+    def extra_repr(self):
+        return f"features={self.features}, patch_size={self.patch_size!r}"
+
+    def forward(self, x, pixel_aspect_ratio=1.0):
+        *_, h, w = x.shape
+        h_out = h // self.patch_size[0]
+        w_out = w // self.patch_size[1]
+        if h % self.patch_size[0] != 0 or w % self.patch_size[1] != 0:
+            raise ValueError(f"Image size {h}x{w} is not divisible by patch size {self.patch_size[0]}x{self.patch_size[1]}")
+        x = rearrange(x, "... c (h i) (w j) -> ... (h w) (c i j)", i=self.patch_size[0], j=self.patch_size[1])
+        pixel_aspect_ratio = pixel_aspect_ratio * self.patch_size[0] / self.patch_size[1]
+        pos = make_axial_pos(h_out, w_out, pixel_aspect_ratio, device=x.device)
+        return x, pos
+
+
+class Unpatching(nn.Module):
+    def __init__(self, features, patch_size):
+        super().__init__()
+        self.features = features
+        self.patch_size = patch_size
+        self.d_in = features * patch_size[0] * patch_size[1]
+
+    def extra_repr(self):
+        return f"features={self.features}, patch_size={self.patch_size!r}"
+
+    def forward(self, x, h, w):
+        h_in = h // self.patch_size[0]
+        w_in = w // self.patch_size[1]
+        x = rearrange(x, "... (h w) (c i j) -> ... c (h i) (w j)", h=h_in, w=w_in, i=self.patch_size[0], j=self.patch_size[1])
+        return x
+
+
+class MappingFeedForwardBlock(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.norm = RMSNorm(d_model)
+        self.up_proj = apply_wd(nn.Linear(d_model, d_ff * 2, bias=False))
+        self.act = GEGLU()
+        self.dropout = nn.Dropout(dropout)
+        self.down_proj = apply_wd(zero_init(nn.Linear(d_ff, d_model, bias=False)))
+
+    def forward(self, x):
+        skip = x
+        x = self.norm(x)
+        x = self.up_proj(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.down_proj(x)
+        return x + skip
+
+
+class MappingNetwork(nn.Module):
+    def __init__(self, n_layers, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.in_norm = RMSNorm(d_model)
+        self.blocks = nn.ModuleList([MappingFeedForwardBlock(d_model, d_ff, dropout=dropout) for _ in range(n_layers)])
+        self.out_norm = RMSNorm(d_model)
+
+    def forward(self, x):
+        x = self.in_norm(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.out_norm(x)
+        return x
+
+
+class ImageTransformerDenoiserModelV1(nn.Module):
+    def __init__(self, n_layers, d_model, d_ff, in_features, out_features, patch_size, num_classes=0, dropout=0.0, sigma_data=1.0):
+        super().__init__()
+        self.sigma_data = sigma_data
+        self.num_classes = num_classes
+        self.patch_in = Patching(in_features, patch_size)
+        self.patch_out = Unpatching(out_features, patch_size)
+
+        self.time_emb = layers.FourierFeatures(1, d_model)
+        self.time_in_proj = nn.Linear(d_model, d_model, bias=False)
+        self.aug_emb = layers.FourierFeatures(9, d_model)
+        self.aug_in_proj = nn.Linear(d_model, d_model, bias=False)
+        self.class_emb = nn.Embedding(num_classes, d_model) if num_classes else None
+        self.mapping = tag_module(MappingNetwork(2, d_model, d_ff, dropout=dropout), "mapping")
+
+        self.in_proj = nn.Linear(self.patch_in.d_out, d_model, bias=False)
+        self.blocks = nn.ModuleList([TransformerBlock(d_model, d_ff, 64, dropout=dropout) for _ in range(n_layers)])
+        self.out_norm = RMSNorm(d_model)
+        self.out_proj = zero_init(nn.Linear(d_model, self.patch_out.d_in, bias=False))
+
+    def proj_(self):
+        for block in self.blocks:
+            block.self_attn.qk_norm.proj_()
+
+    def param_groups(self, base_lr=5e-4, mapping_lr_scale=1 / 3):
+        wd = filter_params(lambda tags: "wd" in tags and "mapping" not in tags, self)
+        no_wd = filter_params(lambda tags: "wd" not in tags and "mapping" not in tags, self)
+        mapping_wd = filter_params(lambda tags: "wd" in tags and "mapping" in tags, self)
+        mapping_no_wd = filter_params(lambda tags: "wd" not in tags and "mapping" in tags, self)
+        groups = [
+            {"params": list(wd), "lr": base_lr},
+            {"params": list(no_wd), "lr": base_lr, "weight_decay": 0.0},
+            {"params": list(mapping_wd), "lr": base_lr * mapping_lr_scale},
+            {"params": list(mapping_no_wd), "lr": base_lr * mapping_lr_scale, "weight_decay": 0.0}
+        ]
+        return groups
+
+    def forward(self, x, sigma, aug_cond=None, class_cond=None):
+        # Patching
+        *_, h, w = x.shape
+        x, pos = self.patch_in(x)
+        attn_mask = None
+        x = self.in_proj(x)
+
+        # Mapping network
+        if class_cond is None and self.class_emb is not None:
+            raise ValueError("class_cond must be specified if num_classes > 0")
+
+        c_noise = torch.log(sigma) / 4
+        time_emb = self.time_in_proj(self.time_emb(c_noise[..., None]))
+        aug_cond = x.new_zeros([x.shape[0], 9]) if aug_cond is None else aug_cond
+        aug_emb = self.aug_in_proj(self.aug_emb(aug_cond))
+        class_emb = self.class_emb(class_cond) if self.class_emb is not None else 0
+        cond = self.mapping(time_emb + aug_emb + class_emb).unsqueeze(-2)
+
+        # Transformer
+        for block in self.blocks:
+            x = block(x, pos, attn_mask, cond)
+
+        # Unpatching
+        x = self.out_norm(x)
+        x = self.out_proj(x)
+        x = self.patch_out(x, h, w)
+
+        return x
diff --git a/k_diffusion/models/image_transformer_v2.py b/k_diffusion/models/image_transformer_v2.py
new file mode 100644
index 0000000..303c91e
--- /dev/null
+++ b/k_diffusion/models/image_transformer_v2.py
@@ -0,0 +1,753 @@
+"""k-diffusion transformer diffusion models, version 2."""
+
+from dataclasses import dataclass
+from functools import lru_cache, reduce
+import math
+from typing import Union
+
+from einops import rearrange
+import torch
+from torch import nn
+import torch._dynamo
+from torch.nn import functional as F
+
+from . import flags, flops
+from .. import layers
+from .axial_rope import make_axial_pos
+
+
+try:
+    import natten
+except ImportError:
+    natten = None
+
+try:
+    import flash_attn
+except ImportError:
+    flash_attn = None
+
+
+if flags.get_use_compile():
+    torch._dynamo.config.cache_size_limit = max(64, torch._dynamo.config.cache_size_limit)
+    torch._dynamo.config.suppress_errors = True
+
+
+# Helpers
+
+def zero_init(layer):
+    nn.init.zeros_(layer.weight)
+    if layer.bias is not None:
+        nn.init.zeros_(layer.bias)
+    return layer
+
+
+def checkpoint(function, *args, **kwargs):
+    if flags.get_checkpointing():
+        kwargs.setdefault("use_reentrant", True)
+        return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+    else:
+        return function(*args, **kwargs)
+
+
+def downscale_pos(pos):
+    pos = rearrange(pos, "... (h nh) (w nw) e -> ... h w (nh nw) e", nh=2, nw=2)
+    return torch.mean(pos, dim=-2)
+
+
+# Param tags
+
+def tag_param(param, tag):
+    if not hasattr(param, "_tags"):
+        param._tags = set([tag])
+    else:
+        param._tags.add(tag)
+    return param
+
+
+def tag_module(module, tag):
+    for param in module.parameters():
+        tag_param(param, tag)
+    return module
+
+
+def apply_wd(module):
+    for name, param in module.named_parameters():
+        if name.endswith("weight"):
+            tag_param(param, "wd")
+    return module
+
+
+def filter_params(function, module):
+    for param in module.parameters():
+        tags = getattr(param, "_tags", set())
+        if function(tags):
+            yield param
+
+
+# Kernels
+
+@flags.compile_wrap
+def linear_geglu(x, weight, bias=None):
+    x = x @ weight.mT
+    if bias is not None:
+        x = x + bias
+    x, gate = x.chunk(2, dim=-1)
+    return x * F.gelu(gate)
+
+
+@flags.compile_wrap
+def rms_norm(x, scale, eps):
+    dtype = reduce(torch.promote_types, (x.dtype, scale.dtype, torch.float32))
+    mean_sq = torch.mean(x.to(dtype)**2, dim=-1, keepdim=True)
+    scale = scale.to(dtype) * torch.rsqrt(mean_sq + eps)
+    return x * scale.to(x.dtype)
+
+
+@flags.compile_wrap
+def scale_for_cosine_sim(q, k, scale, eps):
+    dtype = reduce(torch.promote_types, (q.dtype, k.dtype, scale.dtype, torch.float32))
+    sum_sq_q = torch.sum(q.to(dtype)**2, dim=-1, keepdim=True)
+    sum_sq_k = torch.sum(k.to(dtype)**2, dim=-1, keepdim=True)
+    sqrt_scale = torch.sqrt(scale.to(dtype))
+    scale_q = sqrt_scale * torch.rsqrt(sum_sq_q + eps)
+    scale_k = sqrt_scale * torch.rsqrt(sum_sq_k + eps)
+    return q * scale_q.to(q.dtype), k * scale_k.to(k.dtype)
+
+
+@flags.compile_wrap
+def scale_for_cosine_sim_qkv(qkv, scale, eps):
+    q, k, v = qkv.unbind(2)
+    q, k = scale_for_cosine_sim(q, k, scale[:, None], eps)
+    return torch.stack((q, k, v), dim=2)
+
+
+# Layers
+
+class Linear(nn.Linear):
+    def forward(self, x):
+        flops.op(flops.op_linear, x.shape, self.weight.shape)
+        return super().forward(x)
+
+
+class LinearGEGLU(nn.Linear):
+    def __init__(self, in_features, out_features, bias=True):
+        super().__init__(in_features, out_features * 2, bias=bias)
+        self.out_features = out_features
+
+    def forward(self, x):
+        flops.op(flops.op_linear, x.shape, self.weight.shape)
+        return linear_geglu(x, self.weight, self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, shape, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(shape))
+
+    def extra_repr(self):
+        return f"shape={tuple(self.scale.shape)}, eps={self.eps}"
+
+    def forward(self, x):
+        return rms_norm(x, self.scale, self.eps)
+
+
+class AdaRMSNorm(nn.Module):
+    def __init__(self, features, cond_features, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.linear = apply_wd(zero_init(Linear(cond_features, features, bias=False)))
+        tag_module(self.linear, "mapping")
+
+    def extra_repr(self):
+        return f"eps={self.eps},"
+
+    def forward(self, x, cond):
+        return rms_norm(x, self.linear(cond)[:, None, None, :] + 1, self.eps)
+
+
+# Rotary position embeddings
+
+@flags.compile_wrap
+def apply_rotary_emb(x, theta, conj=False):
+    out_dtype = x.dtype
+    dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32))
+    d = theta.shape[-1]
+    assert d * 2 <= x.shape[-1]
+    x1, x2, x3 = x[..., :d], x[..., d : d * 2], x[..., d * 2 :]
+    x1, x2, theta = x1.to(dtype), x2.to(dtype), theta.to(dtype)
+    cos, sin = torch.cos(theta), torch.sin(theta)
+    sin = -sin if conj else sin
+    y1 = x1 * cos - x2 * sin
+    y2 = x2 * cos + x1 * sin
+    y1, y2 = y1.to(out_dtype), y2.to(out_dtype)
+    return torch.cat((y1, y2, x3), dim=-1)
+
+
+@flags.compile_wrap
+def _apply_rotary_emb_inplace(x, theta, conj):
+    dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32))
+    d = theta.shape[-1]
+    assert d * 2 <= x.shape[-1]
+    x1, x2 = x[..., :d], x[..., d : d * 2]
+    x1_, x2_, theta = x1.to(dtype), x2.to(dtype), theta.to(dtype)
+    cos, sin = torch.cos(theta), torch.sin(theta)
+    sin = -sin if conj else sin
+    y1 = x1_ * cos - x2_ * sin
+    y2 = x2_ * cos + x1_ * sin
+    x1.copy_(y1)
+    x2.copy_(y2)
+
+
+class ApplyRotaryEmbeddingInplace(torch.autograd.Function):
+    @staticmethod
+    def forward(x, theta, conj):
+        _apply_rotary_emb_inplace(x, theta, conj=conj)
+        return x
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        _, theta, conj = inputs
+        ctx.save_for_backward(theta)
+        ctx.conj = conj
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        theta, = ctx.saved_tensors
+        _apply_rotary_emb_inplace(grad_output, theta, conj=not ctx.conj)
+        return grad_output, None, None
+
+
+def apply_rotary_emb_(x, theta):
+    return ApplyRotaryEmbeddingInplace.apply(x, theta, False)
+
+
+class AxialRoPE(nn.Module):
+    def __init__(self, dim, n_heads):
+        super().__init__()
+        log_min = math.log(math.pi)
+        log_max = math.log(10.0 * math.pi)
+        freqs = torch.linspace(log_min, log_max, n_heads * dim // 4 + 1)[:-1].exp()
+        self.register_buffer("freqs", freqs.view(dim // 4, n_heads).T.contiguous())
+
+    def extra_repr(self):
+        return f"dim={self.freqs.shape[1] * 4}, n_heads={self.freqs.shape[0]}"
+
+    def forward(self, pos):
+        theta_h = pos[..., None, 0:1] * self.freqs.to(pos.dtype)
+        theta_w = pos[..., None, 1:2] * self.freqs.to(pos.dtype)
+        return torch.cat((theta_h, theta_w), dim=-1)
+
+
+# Shifted window attention
+
+def window(window_size, x):
+    *b, h, w, c = x.shape
+    x = torch.reshape(
+        x,
+        (*b, h // window_size, window_size, w // window_size, window_size, c),
+    )
+    x = torch.permute(
+        x,
+        (*range(len(b)), -5, -3, -4, -2, -1),
+    )
+    return x
+
+
+def unwindow(x):
+    *b, h, w, wh, ww, c = x.shape
+    x = torch.permute(x, (*range(len(b)), -5, -3, -4, -2, -1))
+    x = torch.reshape(x, (*b, h * wh, w * ww, c))
+    return x
+
+
+def shifted_window(window_size, window_shift, x):
+    x = torch.roll(x, shifts=(window_shift, window_shift), dims=(-2, -3))
+    windows = window(window_size, x)
+    return windows
+
+
+def shifted_unwindow(window_shift, x):
+    x = unwindow(x)
+    x = torch.roll(x, shifts=(-window_shift, -window_shift), dims=(-2, -3))
+    return x
+
+
+@lru_cache
+def make_shifted_window_masks(n_h_w, n_w_w, w_h, w_w, shift, device=None):
+    ph_coords = torch.arange(n_h_w, device=device)
+    pw_coords = torch.arange(n_w_w, device=device)
+    h_coords = torch.arange(w_h, device=device)
+    w_coords = torch.arange(w_w, device=device)
+    patch_h, patch_w, q_h, q_w, k_h, k_w = torch.meshgrid(
+        ph_coords,
+        pw_coords,
+        h_coords,
+        w_coords,
+        h_coords,
+        w_coords,
+        indexing="ij",
+    )
+    is_top_patch = patch_h == 0
+    is_left_patch = patch_w == 0
+    q_above_shift = q_h < shift
+    k_above_shift = k_h < shift
+    q_left_of_shift = q_w < shift
+    k_left_of_shift = k_w < shift
+    m_corner = (
+        is_left_patch
+        & is_top_patch
+        & (q_left_of_shift == k_left_of_shift)
+        & (q_above_shift == k_above_shift)
+    )
+    m_left = is_left_patch & ~is_top_patch & (q_left_of_shift == k_left_of_shift)
+    m_top = ~is_left_patch & is_top_patch & (q_above_shift == k_above_shift)
+    m_rest = ~is_left_patch & ~is_top_patch
+    m = m_corner | m_left | m_top | m_rest
+    return m
+
+
+def apply_window_attention(window_size, window_shift, q, k, v, scale=None):
+    # prep windows and masks
+    q_windows = shifted_window(window_size, window_shift, q)
+    k_windows = shifted_window(window_size, window_shift, k)
+    v_windows = shifted_window(window_size, window_shift, v)
+    b, heads, h, w, wh, ww, d_head = q_windows.shape
+    mask = make_shifted_window_masks(h, w, wh, ww, window_shift, device=q.device)
+    q_seqs = torch.reshape(q_windows, (b, heads, h, w, wh * ww, d_head))
+    k_seqs = torch.reshape(k_windows, (b, heads, h, w, wh * ww, d_head))
+    v_seqs = torch.reshape(v_windows, (b, heads, h, w, wh * ww, d_head))
+    mask = torch.reshape(mask, (h, w, wh * ww, wh * ww))
+
+    # do the attention here
+    flops.op(flops.op_attention, q_seqs.shape, k_seqs.shape, v_seqs.shape)
+    qkv = F.scaled_dot_product_attention(q_seqs, k_seqs, v_seqs, mask, scale=scale)
+
+    # unwindow
+    qkv = torch.reshape(qkv, (b, heads, h, w, wh, ww, d_head))
+    return shifted_unwindow(window_shift, qkv)
+
+
+# Transformer layers
+
+
+def use_flash_2(x):
+    if not flags.get_use_flash_attention_2():
+        return False
+    if flash_attn is None:
+        return False
+    if x.device.type != "cuda":
+        return False
+    if x.dtype not in (torch.float16, torch.bfloat16):
+        return False
+    return True
+
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, cond_features, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False))
+        self.scale = nn.Parameter(torch.full([self.n_heads], 10.0))
+        self.pos_emb = AxialRoPE(d_head // 2, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head},"
+
+    def forward(self, x, pos, cond):
+        skip = x
+        x = self.norm(x, cond)
+        qkv = self.qkv_proj(x)
+        pos = rearrange(pos, "... h w e -> ... (h w) e").to(qkv.dtype)
+        theta = self.pos_emb(pos)
+        if use_flash_2(qkv):
+            qkv = rearrange(qkv, "n h w (t nh e) -> n (h w) t nh e", t=3, e=self.d_head)
+            qkv = scale_for_cosine_sim_qkv(qkv, self.scale, 1e-6)
+            theta = torch.stack((theta, theta, torch.zeros_like(theta)), dim=-3)
+            qkv = apply_rotary_emb_(qkv, theta)
+            flops_shape = qkv.shape[-5], qkv.shape[-2], qkv.shape[-4], qkv.shape[-1]
+            flops.op(flops.op_attention, flops_shape, flops_shape, flops_shape)
+            x = flash_attn.flash_attn_qkvpacked_func(qkv, softmax_scale=1.0)
+            x = rearrange(x, "n (h w) nh e -> n h w (nh e)", h=skip.shape[-3], w=skip.shape[-2])
+        else:
+            q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh (h w) e", t=3, e=self.d_head)
+            q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None], 1e-6)
+            theta = theta.movedim(-2, -3)
+            q = apply_rotary_emb_(q, theta)
+            k = apply_rotary_emb_(k, theta)
+            flops.op(flops.op_attention, q.shape, k.shape, v.shape)
+            x = F.scaled_dot_product_attention(q, k, v, scale=1.0)
+            x = rearrange(x, "n nh (h w) e -> n h w (nh e)", h=skip.shape[-3], w=skip.shape[-2])
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class NeighborhoodSelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, cond_features, kernel_size, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.kernel_size = kernel_size
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False))
+        self.scale = nn.Parameter(torch.full([self.n_heads], 10.0))
+        self.pos_emb = AxialRoPE(d_head // 2, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head}, kernel_size={self.kernel_size}"
+
+    def forward(self, x, pos, cond):
+        skip = x
+        x = self.norm(x, cond)
+        qkv = self.qkv_proj(x)
+        if natten is None:
+            raise ModuleNotFoundError("natten is required for neighborhood attention")
+        if natten.has_fused_na():
+            q, k, v = rearrange(qkv, "n h w (t nh e) -> t n h w nh e", t=3, e=self.d_head)
+            q, k = scale_for_cosine_sim(q, k, self.scale[:, None], 1e-6)
+            theta = self.pos_emb(pos)
+            q = apply_rotary_emb_(q, theta)
+            k = apply_rotary_emb_(k, theta)
+            flops.op(flops.op_natten, q.shape, k.shape, v.shape, self.kernel_size)
+            x = natten.functional.na2d(q, k, v, self.kernel_size, scale=1.0)
+            x = rearrange(x, "n h w nh e -> n h w (nh e)")
+        else:
+            q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh h w e", t=3, e=self.d_head)
+            q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None, None], 1e-6)
+            theta = self.pos_emb(pos).movedim(-2, -4)
+            q = apply_rotary_emb_(q, theta)
+            k = apply_rotary_emb_(k, theta)
+            flops.op(flops.op_natten, q.shape, k.shape, v.shape, self.kernel_size)
+            qk = natten.functional.na2d_qk(q, k, self.kernel_size)
+            a = torch.softmax(qk, dim=-1).to(v.dtype)
+            x = natten.functional.na2d_av(a, v, self.kernel_size)
+            x = rearrange(x, "n nh h w e -> n h w (nh e)")
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class ShiftedWindowSelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, cond_features, window_size, window_shift, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.window_size = window_size
+        self.window_shift = window_shift
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False))
+        self.scale = nn.Parameter(torch.full([self.n_heads], 10.0))
+        self.pos_emb = AxialRoPE(d_head // 2, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head}, window_size={self.window_size}, window_shift={self.window_shift}"
+
+    def forward(self, x, pos, cond):
+        skip = x
+        x = self.norm(x, cond)
+        qkv = self.qkv_proj(x)
+        q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh h w e", t=3, e=self.d_head)
+        q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None, None], 1e-6)
+        theta = self.pos_emb(pos).movedim(-2, -4)
+        q = apply_rotary_emb_(q, theta)
+        k = apply_rotary_emb_(k, theta)
+        x = apply_window_attention(self.window_size, self.window_shift, q, k, v, scale=1.0)
+        x = rearrange(x, "n nh h w e -> n h w (nh e)")
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class FeedForwardBlock(nn.Module):
+    def __init__(self, d_model, d_ff, cond_features, dropout=0.0):
+        super().__init__()
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.up_proj = apply_wd(LinearGEGLU(d_model, d_ff, bias=False))
+        self.dropout = nn.Dropout(dropout)
+        self.down_proj = apply_wd(zero_init(Linear(d_ff, d_model, bias=False)))
+
+    def forward(self, x, cond):
+        skip = x
+        x = self.norm(x, cond)
+        x = self.up_proj(x)
+        x = self.dropout(x)
+        x = self.down_proj(x)
+        return x + skip
+
+
+class GlobalTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, cond_features, dropout=0.0):
+        super().__init__()
+        self.self_attn = SelfAttentionBlock(d_model, d_head, cond_features, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.self_attn, x, pos, cond)
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class NeighborhoodTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, cond_features, kernel_size, dropout=0.0):
+        super().__init__()
+        self.self_attn = NeighborhoodSelfAttentionBlock(d_model, d_head, cond_features, kernel_size, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.self_attn, x, pos, cond)
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class ShiftedWindowTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, cond_features, window_size, index, dropout=0.0):
+        super().__init__()
+        window_shift = window_size // 2 if index % 2 == 1 else 0
+        self.self_attn = ShiftedWindowSelfAttentionBlock(d_model, d_head, cond_features, window_size, window_shift, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.self_attn, x, pos, cond)
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class NoAttentionTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, cond_features, dropout=0.0):
+        super().__init__()
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class Level(nn.ModuleList):
+    def forward(self, x, *args, **kwargs):
+        for layer in self:
+            x = layer(x, *args, **kwargs)
+        return x
+
+
+# Mapping network
+
+class MappingFeedForwardBlock(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.norm = RMSNorm(d_model)
+        self.up_proj = apply_wd(LinearGEGLU(d_model, d_ff, bias=False))
+        self.dropout = nn.Dropout(dropout)
+        self.down_proj = apply_wd(zero_init(Linear(d_ff, d_model, bias=False)))
+
+    def forward(self, x):
+        skip = x
+        x = self.norm(x)
+        x = self.up_proj(x)
+        x = self.dropout(x)
+        x = self.down_proj(x)
+        return x + skip
+
+
+class MappingNetwork(nn.Module):
+    def __init__(self, n_layers, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.in_norm = RMSNorm(d_model)
+        self.blocks = nn.ModuleList([MappingFeedForwardBlock(d_model, d_ff, dropout=dropout) for _ in range(n_layers)])
+        self.out_norm = RMSNorm(d_model)
+
+    def forward(self, x):
+        x = self.in_norm(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.out_norm(x)
+        return x
+
+
+# Token merging and splitting
+
+class TokenMerge(nn.Module):
+    def __init__(self, in_features, out_features, patch_size=(2, 2)):
+        super().__init__()
+        self.h = patch_size[0]
+        self.w = patch_size[1]
+        self.proj = apply_wd(Linear(in_features * self.h * self.w, out_features, bias=False))
+
+    def forward(self, x):
+        x = rearrange(x, "... (h nh) (w nw) e -> ... h w (nh nw e)", nh=self.h, nw=self.w)
+        return self.proj(x)
+
+
+class TokenSplitWithoutSkip(nn.Module):
+    def __init__(self, in_features, out_features, patch_size=(2, 2)):
+        super().__init__()
+        self.h = patch_size[0]
+        self.w = patch_size[1]
+        self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False))
+
+    def forward(self, x):
+        x = self.proj(x)
+        return rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w)
+
+
+class TokenSplit(nn.Module):
+    def __init__(self, in_features, out_features, patch_size=(2, 2)):
+        super().__init__()
+        self.h = patch_size[0]
+        self.w = patch_size[1]
+        self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False))
+        self.fac = nn.Parameter(torch.ones(1) * 0.5)
+
+    def forward(self, x, skip):
+        x = self.proj(x)
+        x = rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w)
+        return torch.lerp(skip, x, self.fac.to(x.dtype))
+
+
+# Configuration
+
+@dataclass
+class GlobalAttentionSpec:
+    d_head: int
+
+
+@dataclass
+class NeighborhoodAttentionSpec:
+    d_head: int
+    kernel_size: int
+
+
+@dataclass
+class ShiftedWindowAttentionSpec:
+    d_head: int
+    window_size: int
+
+
+@dataclass
+class NoAttentionSpec:
+    pass
+
+
+@dataclass
+class LevelSpec:
+    depth: int
+    width: int
+    d_ff: int
+    self_attn: Union[GlobalAttentionSpec, NeighborhoodAttentionSpec, ShiftedWindowAttentionSpec, NoAttentionSpec]
+    dropout: float
+
+
+@dataclass
+class MappingSpec:
+    depth: int
+    width: int
+    d_ff: int
+    dropout: float
+
+
+# Model class
+
+class ImageTransformerDenoiserModelV2(nn.Module):
+    def __init__(self, levels, mapping, in_channels, out_channels, patch_size, num_classes=0, mapping_cond_dim=0):
+        super().__init__()
+        self.num_classes = num_classes
+
+        self.patch_in = TokenMerge(in_channels, levels[0].width, patch_size)
+
+        self.time_emb = layers.FourierFeatures(1, mapping.width)
+        self.time_in_proj = Linear(mapping.width, mapping.width, bias=False)
+        self.aug_emb = layers.FourierFeatures(9, mapping.width)
+        self.aug_in_proj = Linear(mapping.width, mapping.width, bias=False)
+        self.class_emb = nn.Embedding(num_classes, mapping.width) if num_classes else None
+        self.mapping_cond_in_proj = Linear(mapping_cond_dim, mapping.width, bias=False) if mapping_cond_dim else None
+        self.mapping = tag_module(MappingNetwork(mapping.depth, mapping.width, mapping.d_ff, dropout=mapping.dropout), "mapping")
+
+        self.down_levels, self.up_levels = nn.ModuleList(), nn.ModuleList()
+        for i, spec in enumerate(levels):
+            if isinstance(spec.self_attn, GlobalAttentionSpec):
+                layer_factory = lambda _: GlobalTransformerLayer(spec.width, spec.d_ff, spec.self_attn.d_head, mapping.width, dropout=spec.dropout)
+            elif isinstance(spec.self_attn, NeighborhoodAttentionSpec):
+                layer_factory = lambda _: NeighborhoodTransformerLayer(spec.width, spec.d_ff, spec.self_attn.d_head, mapping.width, spec.self_attn.kernel_size, dropout=spec.dropout)
+            elif isinstance(spec.self_attn, ShiftedWindowAttentionSpec):
+                layer_factory = lambda i: ShiftedWindowTransformerLayer(spec.width, spec.d_ff, spec.self_attn.d_head, mapping.width, spec.self_attn.window_size, i, dropout=spec.dropout)
+            elif isinstance(spec.self_attn, NoAttentionSpec):
+                layer_factory = lambda _: NoAttentionTransformerLayer(spec.width, spec.d_ff, mapping.width, dropout=spec.dropout)
+            else:
+                raise ValueError(f"unsupported self attention spec {spec.self_attn}")
+
+            if i < len(levels) - 1:
+                self.down_levels.append(Level([layer_factory(i) for i in range(spec.depth)]))
+                self.up_levels.append(Level([layer_factory(i + spec.depth) for i in range(spec.depth)]))
+            else:
+                self.mid_level = Level([layer_factory(i) for i in range(spec.depth)])
+
+        self.merges = nn.ModuleList([TokenMerge(spec_1.width, spec_2.width) for spec_1, spec_2 in zip(levels[:-1], levels[1:])])
+        self.splits = nn.ModuleList([TokenSplit(spec_2.width, spec_1.width) for spec_1, spec_2 in zip(levels[:-1], levels[1:])])
+
+        self.out_norm = RMSNorm(levels[0].width)
+        self.patch_out = TokenSplitWithoutSkip(levels[0].width, out_channels, patch_size)
+        nn.init.zeros_(self.patch_out.proj.weight)
+
+    def param_groups(self, base_lr=5e-4, mapping_lr_scale=1 / 3):
+        wd = filter_params(lambda tags: "wd" in tags and "mapping" not in tags, self)
+        no_wd = filter_params(lambda tags: "wd" not in tags and "mapping" not in tags, self)
+        mapping_wd = filter_params(lambda tags: "wd" in tags and "mapping" in tags, self)
+        mapping_no_wd = filter_params(lambda tags: "wd" not in tags and "mapping" in tags, self)
+        groups = [
+            {"params": list(wd), "lr": base_lr},
+            {"params": list(no_wd), "lr": base_lr, "weight_decay": 0.0},
+            {"params": list(mapping_wd), "lr": base_lr * mapping_lr_scale},
+            {"params": list(mapping_no_wd), "lr": base_lr * mapping_lr_scale, "weight_decay": 0.0}
+        ]
+        return groups
+
+    def forward(self, x, sigma, aug_cond=None, class_cond=None, mapping_cond=None):
+        # Patching
+        x = x.movedim(-3, -1)
+        x = self.patch_in(x)
+        # TODO: pixel aspect ratio for nonsquare patches
+        pos = make_axial_pos(x.shape[-3], x.shape[-2], device=x.device).view(x.shape[-3], x.shape[-2], 2)
+
+        # Mapping network
+        if class_cond is None and self.class_emb is not None:
+            raise ValueError("class_cond must be specified if num_classes > 0")
+        if mapping_cond is None and self.mapping_cond_in_proj is not None:
+            raise ValueError("mapping_cond must be specified if mapping_cond_dim > 0")
+
+        c_noise = torch.log(sigma) / 4
+        time_emb = self.time_in_proj(self.time_emb(c_noise[..., None]))
+        aug_cond = x.new_zeros([x.shape[0], 9]) if aug_cond is None else aug_cond
+        aug_emb = self.aug_in_proj(self.aug_emb(aug_cond))
+        class_emb = self.class_emb(class_cond) if self.class_emb is not None else 0
+        mapping_emb = self.mapping_cond_in_proj(mapping_cond) if self.mapping_cond_in_proj is not None else 0
+        cond = self.mapping(time_emb + aug_emb + class_emb + mapping_emb)
+
+        # Hourglass transformer
+        skips, poses = [], []
+        for down_level, merge in zip(self.down_levels, self.merges):
+            x = down_level(x, pos, cond)
+            skips.append(x)
+            poses.append(pos)
+            x = merge(x)
+            pos = downscale_pos(pos)
+
+        x = self.mid_level(x, pos, cond)
+
+        for up_level, split, skip, pos in reversed(list(zip(self.up_levels, self.splits, skips, poses))):
+            x = split(x, skip)
+            x = up_level(x, pos, cond)
+
+        # Unpatching
+        x = self.out_norm(x)
+        x = self.patch_out(x)
+        x = x.movedim(-1, -3)
+
+        return x
diff --git a/k_diffusion/models/image_v1.py b/k_diffusion/models/image_v1.py
new file mode 100644
index 0000000..88a1145
--- /dev/null
+++ b/k_diffusion/models/image_v1.py
@@ -0,0 +1,176 @@
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .. import layers, utils
+
+
+def orthogonal_(module):
+    nn.init.orthogonal_(module.weight)
+    return module
+
+
+class ResConvBlock(layers.ConditionedResidualBlock):
+    def __init__(self, feats_in, c_in, c_mid, c_out, group_size=32, dropout_rate=0.):
+        skip = None if c_in == c_out else orthogonal_(nn.Conv2d(c_in, c_out, 1, bias=False))
+        super().__init__(
+            layers.AdaGN(feats_in, c_in, max(1, c_in // group_size)),
+            nn.GELU(),
+            nn.Conv2d(c_in, c_mid, 3, padding=1),
+            nn.Dropout2d(dropout_rate, inplace=True),
+            layers.AdaGN(feats_in, c_mid, max(1, c_mid // group_size)),
+            nn.GELU(),
+            nn.Conv2d(c_mid, c_out, 3, padding=1),
+            nn.Dropout2d(dropout_rate, inplace=True),
+            skip=skip)
+        nn.init.zeros_(self.main[-2].weight)
+        nn.init.zeros_(self.main[-2].bias)
+
+
+class DBlock(layers.ConditionedSequential):
+    def __init__(self, n_layers, feats_in, c_in, c_mid, c_out, group_size=32, head_size=64, dropout_rate=0., downsample=False, self_attn=False, cross_attn=False, c_enc=0):
+        modules = [nn.Identity()]
+        for i in range(n_layers):
+            my_c_in = c_in if i == 0 else c_mid
+            my_c_out = c_mid if i < n_layers - 1 else c_out
+            modules.append(ResConvBlock(feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate))
+            if self_attn:
+                norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
+                modules.append(layers.SelfAttention2d(my_c_out, max(1, my_c_out // head_size), norm, dropout_rate))
+            if cross_attn:
+                norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
+                modules.append(layers.CrossAttention2d(my_c_out, c_enc, max(1, my_c_out // head_size), norm, dropout_rate))
+        super().__init__(*modules)
+        self.set_downsample(downsample)
+
+    def set_downsample(self, downsample):
+        self[0] = layers.Downsample2d() if downsample else nn.Identity()
+        return self
+
+
+class UBlock(layers.ConditionedSequential):
+    def __init__(self, n_layers, feats_in, c_in, c_mid, c_out, group_size=32, head_size=64, dropout_rate=0., upsample=False, self_attn=False, cross_attn=False, c_enc=0):
+        modules = []
+        for i in range(n_layers):
+            my_c_in = c_in if i == 0 else c_mid
+            my_c_out = c_mid if i < n_layers - 1 else c_out
+            modules.append(ResConvBlock(feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate))
+            if self_attn:
+                norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
+                modules.append(layers.SelfAttention2d(my_c_out, max(1, my_c_out // head_size), norm, dropout_rate))
+            if cross_attn:
+                norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
+                modules.append(layers.CrossAttention2d(my_c_out, c_enc, max(1, my_c_out // head_size), norm, dropout_rate))
+        modules.append(nn.Identity())
+        super().__init__(*modules)
+        self.set_upsample(upsample)
+
+    def forward(self, input, cond, skip=None):
+        if skip is not None:
+            input = torch.cat([input, skip], dim=1)
+        return super().forward(input, cond)
+
+    def set_upsample(self, upsample):
+        self[-1] = layers.Upsample2d() if upsample else nn.Identity()
+        return self
+
+
+class MappingNet(nn.Sequential):
+    def __init__(self, feats_in, feats_out, n_layers=2):
+        layers = []
+        for i in range(n_layers):
+            layers.append(orthogonal_(nn.Linear(feats_in if i == 0 else feats_out, feats_out)))
+            layers.append(nn.GELU())
+        super().__init__(*layers)
+
+
+class ImageDenoiserModelV1(nn.Module):
+    def __init__(self, c_in, feats_in, depths, channels, self_attn_depths, cross_attn_depths=None, mapping_cond_dim=0, unet_cond_dim=0, cross_cond_dim=0, dropout_rate=0., patch_size=1, skip_stages=0, has_variance=False):
+        super().__init__()
+        self.c_in = c_in
+        self.channels = channels
+        self.unet_cond_dim = unet_cond_dim
+        self.patch_size = patch_size
+        self.has_variance = has_variance
+        self.timestep_embed = layers.FourierFeatures(1, feats_in)
+        if mapping_cond_dim > 0:
+            self.mapping_cond = nn.Linear(mapping_cond_dim, feats_in, bias=False)
+        self.mapping = MappingNet(feats_in, feats_in)
+        self.proj_in = nn.Conv2d((c_in + unet_cond_dim) * self.patch_size ** 2, channels[max(0, skip_stages - 1)], 1)
+        self.proj_out = nn.Conv2d(channels[max(0, skip_stages - 1)], c_in * self.patch_size ** 2 + (1 if self.has_variance else 0), 1)
+        nn.init.zeros_(self.proj_out.weight)
+        nn.init.zeros_(self.proj_out.bias)
+        if cross_cond_dim == 0:
+            cross_attn_depths = [False] * len(self_attn_depths)
+        d_blocks, u_blocks = [], []
+        for i in range(len(depths)):
+            my_c_in = channels[max(0, i - 1)]
+            d_blocks.append(DBlock(depths[i], feats_in, my_c_in, channels[i], channels[i], downsample=i > skip_stages, self_attn=self_attn_depths[i], cross_attn=cross_attn_depths[i], c_enc=cross_cond_dim, dropout_rate=dropout_rate))
+        for i in range(len(depths)):
+            my_c_in = channels[i] * 2 if i < len(depths) - 1 else channels[i]
+            my_c_out = channels[max(0, i - 1)]
+            u_blocks.append(UBlock(depths[i], feats_in, my_c_in, channels[i], my_c_out, upsample=i > skip_stages, self_attn=self_attn_depths[i], cross_attn=cross_attn_depths[i], c_enc=cross_cond_dim, dropout_rate=dropout_rate))
+        self.u_net = layers.UNet(d_blocks, reversed(u_blocks), skip_stages=skip_stages)
+
+    def param_groups(self, base_lr=2e-4):
+        wd_names = []
+        for name, _ in self.named_parameters():
+            if name.startswith("mapping") or name.startswith("u_net"):
+                if name.endswith(".weight"):
+                    wd_names.append(name)
+        wd, no_wd = [], []
+        for name, param in self.named_parameters():
+            if name in wd_names:
+                wd.append(param)
+            else:
+                no_wd.append(param)
+        groups = [
+            {"params": wd, "lr": base_lr},
+            {"params": no_wd, "lr": base_lr, "weight_decay": 0.0},
+        ]
+        return groups
+
+    def forward(self, input, sigma, mapping_cond=None, unet_cond=None, cross_cond=None, cross_cond_padding=None, return_variance=False):
+        c_noise = sigma.log() / 4
+        timestep_embed = self.timestep_embed(utils.append_dims(c_noise, 2))
+        mapping_cond_embed = torch.zeros_like(timestep_embed) if mapping_cond is None else self.mapping_cond(mapping_cond)
+        mapping_out = self.mapping(timestep_embed + mapping_cond_embed)
+        cond = {'cond': mapping_out}
+        if unet_cond is not None:
+            input = torch.cat([input, unet_cond], dim=1)
+        if cross_cond is not None:
+            cond['cross'] = cross_cond
+            cond['cross_padding'] = cross_cond_padding
+        if self.patch_size > 1:
+            input = F.pixel_unshuffle(input, self.patch_size)
+        input = self.proj_in(input)
+        input = self.u_net(input, cond)
+        input = self.proj_out(input)
+        if self.has_variance:
+            input, logvar = input[:, :-1], input[:, -1].flatten(1).mean(1)
+        if self.patch_size > 1:
+            input = F.pixel_shuffle(input, self.patch_size)
+        if self.has_variance and return_variance:
+            return input, logvar
+        return input
+
+    def set_skip_stages(self, skip_stages):
+        self.proj_in = nn.Conv2d(self.proj_in.in_channels, self.channels[max(0, skip_stages - 1)], 1)
+        self.proj_out = nn.Conv2d(self.channels[max(0, skip_stages - 1)], self.proj_out.out_channels, 1)
+        nn.init.zeros_(self.proj_out.weight)
+        nn.init.zeros_(self.proj_out.bias)
+        self.u_net.skip_stages = skip_stages
+        for i, block in enumerate(self.u_net.d_blocks):
+            block.set_downsample(i > skip_stages)
+        for i, block in enumerate(reversed(self.u_net.u_blocks)):
+            block.set_upsample(i > skip_stages)
+        return self
+
+    def set_patch_size(self, patch_size):
+        self.patch_size = patch_size
+        self.proj_in = nn.Conv2d((self.c_in + self.unet_cond_dim) * self.patch_size ** 2, self.channels[max(0, self.u_net.skip_stages - 1)], 1)
+        self.proj_out = nn.Conv2d(self.channels[max(0, self.u_net.skip_stages - 1)], self.c_in * self.patch_size ** 2 + (1 if self.has_variance else 0), 1)
+        nn.init.zeros_(self.proj_out.weight)
+        nn.init.zeros_(self.proj_out.bias)