Upload 3 files

config.json

@@ -1,11 +1,11 @@
 {
   "architectures": [
-    "CustomModelForImageDiffusion"
+    "TinyImageGenModelForImageDiffusion"
   ],
   "auto_map": {
-    "AutoConfig": "configuration_custom.CustomConfig",
-    "AutoModel": "modeling_custom.CustomModel",
-    "AutoModelForImageDiffusion": "modeling_custom.CustomModelForImageDiffusion"
+    "AutoConfig": "configuration_tinyimagegen.TinyImageGenConfig",
+    "AutoModel": "modeling_tinyimagegen.TinyImageGenModel",
+    "AutoModelForImageDiffusion": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion"
   },
   "dtype": "float32",
   "head_dim": 24,
@@ -14,7 +14,7 @@
   "in_channels": 3,
   "initializer_range": 0.02,
   "intermediate_size": 160,
-  "model_type": "custom_image_diffusion",
+  "model_type": "tinyimagegen",
   "num_attention_heads": 4,
   "num_hidden_layers": 6,
   "num_key_value_heads": 2,

configuration_tinyimagegen.py

@@ -0,0 +1,51 @@
+from transformers.configuration_utils import PretrainedConfig
+
+class TinyImageGenConfig(PretrainedConfig):
+    model_type = "tinyimagegen"
+
+    def __init__(
+        self,
+        image_size: int = 32,
+        in_channels: int = 3,
+        patch_size: int = 4,
+        hidden_size: int = 32,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 4,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 48,
+        swiglu_interval: int = 3,
+        num_lanes: int = 4,
+        use_xsa: bool = False,
+        use_per_head_gating: bool = False,
+        rope_theta: float = 2500.0,
+        rms_norm_eps: float = 1e-5,
+        initializer_range: float = 0.02,
+        **kwargs,
+    ):
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.swiglu_interval = swiglu_interval
+        self.num_lanes = num_lanes
+        self.use_xsa = use_xsa
+        self.use_per_head_gating = use_per_head_gating
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.head_dim = hidden_size // num_attention_heads
+        self.num_patches_side = image_size // patch_size
+        self.num_patches = self.num_patches_side ** 2
+        self.patch_dim = in_channels * (patch_size ** 2)
+
+        self.auto_map = {
+            "AutoConfig": "configuration_tinyimagegen.TinyImageGenConfig",
+            "AutoModel": "modeling_tinyimagegen.TinyImageGenModel",
+            "AutoModelForImageDiffusion": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion",
+        }
+
+        super().__init__(**kwargs)

modeling_tinyimagegen.py

@@ -0,0 +1,440 @@
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput
+from safetensors.torch import load_file
+import os
+
+try:
+    from .configuration_tinyimagegen import TinyImageGenConfig
+except Exception:
+    from configuration_tinyimagegen import TinyImageGenConfig
+
+@torch.no_grad()
+def get_hadamard_matrix(d: int, dtype=torch.float32) -> torch.Tensor:
+    p2 = 1 << (d - 1).bit_length()
+    eye = torch.eye(p2, dtype=dtype)
+    h = 1
+    out = eye.clone()
+    while h < p2:
+        out = out.view(-1, 2, h)
+        u = out[:, 0, :]
+        v = out[:, 1, :]
+        out = torch.cat((u + v, u - v), dim=-2)
+        out = out.view(p2, p2)
+        h *= 2
+    out = out * (1.0 / math.sqrt(p2))
+    return out[:d, :d].contiguous()
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * norm * self.weight
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 128):
+        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: torch.Tensor, dim: int, max_period: float = 10000.0) -> torch.Tensor:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t_freq = self.timestep_embedding(t * 1000.0, self.frequency_embedding_size)
+        return self.mlp(t_freq)
+
+class RotaryEmbedding2D(nn.Module):
+    def __init__(self, head_dim: int, base: float = 10000.0):
+        super().__init__()
+        self.head_dim = head_dim
+        self.dim_h = 2 * (head_dim // 4)
+        self.dim_w = head_dim - self.dim_h
+        self.base = base
+
+        inv_freq_h = 1.0 / (self.base ** (torch.arange(0, self.dim_h, 2, dtype=torch.float32) / self.dim_h))
+        inv_freq_w = 1.0 / (self.base ** (torch.arange(0, self.dim_w, 2, dtype=torch.float32) / self.dim_w))
+        self.register_buffer("inv_freq_h", inv_freq_h, persistent=False)
+        self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)
+
+    def forward(self, grid_h: int, grid_w: int, device: torch.device, dtype: torch.dtype = torch.float32):
+        t_h = torch.arange(grid_h, device=device, dtype=torch.float32)
+        t_w = torch.arange(grid_w, device=device, dtype=torch.float32)
+
+        freqs_h = torch.outer(t_h, self.inv_freq_h.to(device=device, dtype=torch.float32))
+        freqs_w = torch.outer(t_w, self.inv_freq_w.to(device=device, dtype=torch.float32))
+
+        emb_h = torch.cat((freqs_h, freqs_h), dim=-1)
+        emb_w = torch.cat((freqs_w, freqs_w), dim=-1)
+
+        emb_h_grid = emb_h[:, None, :].expand(-1, grid_w, -1).reshape(grid_h * grid_w, self.dim_h)
+        emb_w_grid = emb_w[None, :, :].expand(grid_h, -1, -1).reshape(grid_h * grid_w, self.dim_w)
+
+        cos_h = emb_h_grid.cos().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
+        sin_h = emb_h_grid.sin().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
+        cos_w = emb_w_grid.cos().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
+        sin_w = emb_w_grid.sin().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
+        return cos_h, sin_h, cos_w, sin_w
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb_2d(q: torch.Tensor, k: torch.Tensor, cos_h: torch.Tensor, sin_h: torch.Tensor, cos_w: torch.Tensor, sin_w: torch.Tensor):
+    d_h = cos_h.shape[-1]
+    qh, qw = q[..., :d_h], q[..., d_h:]
+    kh, kw = k[..., :d_h], k[..., d_h:]
+
+    qh_rot = (qh * cos_h) + (rotate_half(qh) * sin_h)
+    qw_rot = (qw * cos_w) + (rotate_half(qw) * sin_w)
+    kh_rot = (kh * cos_h) + (rotate_half(kh) * sin_h)
+    kw_rot = (kw * cos_w) + (rotate_half(kw) * sin_w)
+
+    return torch.cat([qh_rot, qw_rot], dim=-1), torch.cat([kh_rot, kw_rot], dim=-1)
+
+class HadamardMLP(nn.Module):
+    def __init__(self, config: TinyImageGenConfig):
+        super().__init__()
+        self.dim = config.hidden_size
+        self.scale1 = nn.Parameter(torch.ones(self.dim))
+        self.scale2 = nn.Parameter(torch.ones(self.dim))
+        self.gate = nn.Parameter(torch.ones(self.dim))
+        self.bias = nn.Parameter(torch.zeros(self.dim))
+
+        hadamard_mat = get_hadamard_matrix(self.dim)
+        self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mat = self.hadamard_mat.type_as(x)
+        h = (x * self.scale1) @ mat
+        g = F.silu(x * self.gate)
+        out = ((h * g) @ mat) * self.scale2 + self.bias
+        return out
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, config: TinyImageGenConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+class XSAGQAttention(nn.Module):
+    def __init__(self, config: TinyImageGenConfig):
+        super().__init__()
+        self.dim = config.hidden_size
+        self.n_heads = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.n_heads // self.n_kv_heads
+        self.use_xsa = config.use_xsa
+        self.use_per_head_gating = config.use_per_head_gating
+
+        self.wq = nn.Linear(self.dim, self.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False)
+
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+        if self.use_per_head_gating:
+            self.head_gate = nn.Linear(self.dim, self.n_heads, bias=True)
+            nn.init.constant_(self.head_gate.bias, 1.0)
+            nn.init.zeros_(self.head_gate.weight)
+
+    def forward(self, x: torch.Tensor, cos_h: torch.Tensor, sin_h: torch.Tensor, cos_w: torch.Tensor, sin_w: torch.Tensor) -> torch.Tensor:
+        bsz, seqlen, _ = x.shape
+
+        xq = self.wq(x).view(bsz, seqlen, self.n_heads, self.head_dim).transpose(1, 2)
+        xk = self.wk(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        xv = self.wv(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        xq = self.q_norm(xq)
+        xk = self.k_norm(xk)
+
+        xq, xk = apply_rotary_pos_emb_2d(xq, xk, cos_h, sin_h, cos_w, sin_w)
+
+        if self.num_kv_groups > 1:
+            xk = xk.repeat_interleave(self.num_kv_groups, dim=1)
+            xv_expanded = xv.repeat_interleave(self.num_kv_groups, dim=1)
+        else:
+            xv_expanded = xv
+
+        attn_out = F.scaled_dot_product_attention(xq, xk, xv_expanded, is_causal=False)
+
+        if self.use_xsa:
+            vn = F.normalize(xv_expanded, p=2, dim=-1, eps=1e-6)
+            proj = (attn_out * vn).sum(dim=-1, keepdim=True)
+            attn_out = attn_out - proj * vn
+
+        if self.use_per_head_gating:
+            gate = torch.sigmoid(self.head_gate(x)).transpose(1, 2).unsqueeze(-1)
+            attn_out = attn_out * gate
+
+        out = attn_out.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        return self.wo(out)
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+class MultiLaneBlock(nn.Module):
+    def __init__(self, config: TinyImageGenConfig, layer_idx: int):
+        super().__init__()
+        self.num_lanes = config.num_lanes
+        self.dim = config.hidden_size
+        self.layer_idx = layer_idx
+
+        self.attn_norm = RMSNorm(self.dim, eps=config.rms_norm_eps)
+        self.attn = XSAGQAttention(config)
+
+        self.mlp_norm = RMSNorm(self.dim, eps=config.rms_norm_eps)
+        if config.swiglu_interval == 0:
+            self.use_swiglu = False
+        elif config.swiglu_interval == 1:
+            self.use_swiglu = True
+        else:
+            self.use_swiglu = ((layer_idx + 1) % config.swiglu_interval == 0)
+
+        if self.use_swiglu:
+            self.mlp = SwiGLUMLP(config)
+        else:
+            self.mlp = HadamardMLP(config)
+
+        self.lane_mix_attn = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes))
+        self.lane_mix_mlp = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes))
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(config.hidden_size, 6 * config.hidden_size, bias=True)
+        )
+        nn.init.zeros_(self.adaLN_modulation[-1].weight)
+        nn.init.zeros_(self.adaLN_modulation[-1].bias)
+
+    def forward(self, lanes: torch.Tensor, t_emb: torch.Tensor, cos_h: torch.Tensor, sin_h: torch.Tensor, cos_w: torch.Tensor, sin_w: torch.Tensor) -> torch.Tensor:
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t_emb).chunk(6, dim=-1)
+
+        primary = lanes[0]
+        normed_primary = modulate(self.attn_norm(primary), shift_msa, scale_msa)
+        attn_update = self.attn(normed_primary, cos_h, sin_h, cos_w, sin_w) * gate_msa.unsqueeze(1)
+
+        mixed = torch.matmul(self.lane_mix_attn, lanes.view(self.num_lanes, -1)).view_as(lanes)
+        lanes = torch.cat([(mixed[0] + attn_update).unsqueeze(0), mixed[1:]], dim=0)
+
+        normed_primary = modulate(self.mlp_norm(lanes[0]), shift_mlp, scale_mlp)
+        mlp_update = self.mlp(normed_primary) * gate_mlp.unsqueeze(1)
+
+        mixed = torch.matmul(self.lane_mix_mlp, lanes.view(self.num_lanes, -1)).view_as(lanes)
+        lanes = torch.cat([(mixed[0] + mlp_update).unsqueeze(0), mixed[1:]], dim=0)
+        return lanes
+
+@dataclass
+class DiffusionOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    v_pred: Optional[torch.FloatTensor] = None
+
+class TinyImageGenPreTrainedModel(PreTrainedModel):
+    config_class = TinyImageGenConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MultiLaneBlock"]
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if hasattr(module, "bias") and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, RMSNorm):
+            module.weight.data.fill_(1.0)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        kwargs.pop("trust_remote_code", None)
+        torch_dtype = kwargs.pop("torch_dtype", None)
+        kwargs.pop("device_map", None)
+        kwargs.pop("low_cpu_mem_usage", None)
+
+        if config is None:
+            config = TinyImageGenConfig.from_pretrained(pretrained_model_name_or_path)
+
+        model = cls(config, *model_args)
+
+        st_file = os.path.join(pretrained_model_name_or_path, "model.safetensors")
+        bin_file = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+
+        if os.path.exists(st_file):
+            state_dict = load_file(st_file)
+            model.load_state_dict(state_dict, strict=False)
+        elif os.path.exists(bin_file):
+            state_dict = torch.load(bin_file, map_location="cpu")
+            model.load_state_dict(state_dict, strict=False)
+        else:
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+
+        if torch_dtype is not None:
+            model.to(dtype=torch_dtype)
+
+        return model
+
+class TinyImageGenModel(TinyImageGenPreTrainedModel):
+    def __init__(self, config: TinyImageGenConfig, *args, **kwargs):
+        super().__init__(config)
+        self.config = config
+        self.num_lanes = config.num_lanes
+        self.gradient_checkpointing = False
+
+        self.x_embedder = nn.Linear(config.patch_dim, config.hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(config.hidden_size)
+        self.rotary_emb = RotaryEmbedding2D(config.head_dim, base=config.rope_theta)
+
+        self.layers = nn.ModuleList([
+            MultiLaneBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)
+        ])
+
+        self.lane_pool_weights = nn.Parameter(torch.tensor([1.0] + [0.1] * (config.num_lanes - 1)))
+        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.final_adaLN = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(config.hidden_size, 2 * config.hidden_size, bias=True)
+        )
+        self.final_proj = nn.Linear(config.hidden_size, config.patch_dim, bias=True)
+
+        nn.init.zeros_(self.final_adaLN[-1].weight)
+        nn.init.zeros_(self.final_adaLN[-1].bias)
+        nn.init.zeros_(self.final_proj.weight)
+        nn.init.zeros_(self.final_proj.bias)
+
+        self.post_init()
+
+    def patchify(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = x.shape
+        p = self.config.patch_size
+        h_patches, w_patches = H // p, W // p
+        x = x.view(B, C, h_patches, p, w_patches, p)
+        x = torch.einsum("bchpwq->bhwpcq", x)
+        x = x.reshape(B, h_patches * w_patches, p * p * C)
+        return x
+
+    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, _ = x.shape
+        p = self.config.patch_size
+        h_patches = self.config.num_patches_side
+        w_patches = self.config.num_patches_side
+        c = self.config.in_channels
+        x = x.reshape(B, h_patches, w_patches, p, p, c)
+        x = torch.einsum("bhwpqc->bchpwq", x)
+        x = x.reshape(B, c, h_patches * p, w_patches * p)
+        return x
+
+    def forward(self, x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        bsz = x_t.shape[0]
+
+        h0 = self.x_embedder(self.patchify(x_t))
+        t_emb = self.t_embedder(t)
+
+        lanes = h0.unsqueeze(0).repeat(self.num_lanes, 1, 1, 1)
+
+        cos_h, sin_h, cos_w, sin_w = self.rotary_emb(self.config.num_patches_side, self.config.num_patches_side, device=x_t.device, dtype=x_t.dtype)
+
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                lanes = cp.checkpoint(layer, lanes, t_emb, cos_h, sin_h, cos_w, sin_w, use_reentrant=False)
+            else:
+                lanes = layer(lanes, t_emb, cos_h, sin_h, cos_w, sin_w)
+
+        pool_weights = F.softmax(self.lane_pool_weights, dim=0).view(self.num_lanes, 1, 1, 1)
+        pooled = (lanes * pool_weights).sum(dim=0)
+
+        shift, scale = self.final_adaLN(t_emb).chunk(2, dim=-1)
+        out = modulate(self.final_norm(pooled), shift, scale)
+        out = self.final_proj(out)
+        return self.unpatchify(out)
+
+    @torch.no_grad()
+    def sample(self, num_samples: int, device: torch.device, num_steps: int = 25) -> torch.Tensor:
+        was_training = self.training
+        self.eval()
+        x = torch.randn((num_samples, self.config.in_channels, self.config.image_size, self.config.image_size), device=device)
+        dt = 1.0 / num_steps
+
+        for step in range(num_steps):
+            t_val = step / num_steps
+            t = torch.full((num_samples,), t_val, device=device, dtype=torch.float32)
+            v = self(x, t)
+            x = x + v * dt
+
+        if was_training:
+            self.train()
+        return x.clamp(-1.0, 1.0)
+
+class TinyImageGenModelForImageDiffusion(TinyImageGenPreTrainedModel):
+    def __init__(self, config: TinyImageGenConfig, *args, **kwargs):
+        super().__init__(config)
+        self.model = TinyImageGenModel(config)
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        x_t: Optional[torch.Tensor] = None,
+        t: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> DiffusionOutput:
+        return_dict = return_dict if return_dict is not None else getattr(self.config, "return_dict", True)
+
+        loss = None
+        v_pred = None
+
+        if pixel_values is not None:
+            x_1 = pixel_values
+            bsz = x_1.shape[0]
+            x_0 = torch.randn_like(x_1)
+            t_rand = torch.rand(bsz, device=x_1.device)
+            t_exp = t_rand.view(bsz, 1, 1, 1)
+
+            x_t_flow = (1.0 - t_exp) * x_0 + t_exp * x_1
+            v_target = x_1 - x_0
+
+            v_pred = self.model(x_t_flow, t_rand)
+            loss = F.mse_loss(v_pred, v_target)
+        elif x_t is not None and t is not None:
+            v_pred = self.model(x_t, t)
+        else:
+            raise ValueError("You must pass either 'pixel_values' for training or ('x_t', 't') for inference.")
+
+        if not return_dict:
+            return (loss, v_pred) if loss is not None else (v_pred,)
+
+        return DiffusionOutput(loss=loss, v_pred=v_pred)
+
+    @torch.no_grad()
+    def sample(self, num_samples: int, device: torch.device, num_steps: int = 25) -> torch.Tensor:
+        return self.model.sample(num_samples=num_samples, device=device, num_steps=num_steps)