Upload full repo excluding dump_40, dump_100, precomputed_tokens, precomputed_data

fb0011a verified 8 months ago

9.29 kB

import math
from typing import Optional

import torch
import torch.nn as nn
from torch import Tensor


# -----------------------
# Positional encodings
# -----------------------
def sinusoidal_positional_encoding(
    length: int, dim: int, device: torch.device
) -> Tensor:
    assert dim % 2 == 0, "positional dim must be even"
    pos = torch.arange(length, device=device).float().unsqueeze(1)  # (L, 1)
    half = dim // 2
    div = torch.exp(
        torch.arange(0, half, device=device).float()
        * (-math.log(10000.0) / max(1, half - 1))
    )  # (half,)
    pe = torch.cat([torch.sin(pos * div), torch.cos(pos * div)], dim=1)
    return pe.unsqueeze(0)  # (1, L, dim)


class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        assert dim % 2 == 0
        self.dim = dim

    def forward(self, t: Tensor) -> Tensor:
        if t.dim() == 2 and t.size(1) == 1:
            t = t[:, 0]
        half = self.dim // 2
        device = t.device
        freq = torch.exp(
            -math.log(10000)
            * torch.arange(half, device=device).float()
            / (half - 1)
        )
        args = t.float().unsqueeze(1) * freq.unsqueeze(0)
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
        return emb  # (B, dim)


# -----------------------
# Transformer building blocks
# -----------------------
class TransformerLayer(nn.Module):
    """
    Minimal layer:
      - self-attn on x
      - cross-attn x -> cond (if provided)
      - MLP
    All with pre-norm and residuals.
    """

    def __init__(self, d_model: int, n_heads: int = 8, dropout: float = 0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.self_attn = nn.MultiheadAttention(
            embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
        )
        self.drop1 = nn.Dropout(dropout)

        self.norm2 = nn.LayerNorm(d_model)
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
        )
        self.drop2 = nn.Dropout(dropout)

        self.norm3 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 4, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x: Tensor, cond: Optional[Tensor]) -> Tensor:
        # x: (B, L, d), cond: (B, T, d) or None
        y = self.norm1(x)
        y, _ = self.self_attn(y, y, y, need_weights=False)
        x = x + self.drop1(y)

        if cond is not None:
            z = self.norm2(x)
            z, _ = self.cross_attn(z, cond, cond, need_weights=False)
            x = x + self.drop2(z)

        w = self.norm3(x)
        x = x + self.mlp(w)
        return x


# -----------------------
# Just-a-Transformer "UNet1D" (no up/down), 3 layers
# Keeps UNet-like forward signature incl. CFG knobs.
# -----------------------
class UNet1DTransformerSimple(nn.Module):
    def __init__(
        self,
        seq_len: int = 64,
        in_channels: int = 1,
        out_channels: int = 1,
        d_model: int = 256,
        n_heads: int = 8,
        dropout: float = 0.0,
        cond_dim: int = 128,  # embedding is [B, T, 128]
        time_emb_dim: int = 128,
    ):
        super().__init__()
        self.seq_len = seq_len
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.d_model = d_model
        self.cond_dim = cond_dim

        # Input/output projections
        self.in_proj = nn.Linear(in_channels, d_model)
        self.out_proj = nn.Linear(d_model, out_channels)

        # Time embedding -> d_model bias
        self.time_pos_emb = SinusoidalPosEmb(time_emb_dim)
        self.time_mlp = nn.Sequential(
            nn.Linear(time_emb_dim, d_model * 4),
            nn.SiLU(),
            nn.Linear(d_model * 4, d_model),
        )

        # Cond projection to d_model (for cross-attention memory)
        self.cond_proj = nn.Linear(cond_dim, d_model)
        # Learnable null cond for classifier-free guidance
        self.null_cond = nn.Parameter(torch.zeros(cond_dim))

        # Exactly 3 transformer layers
        self.layers = nn.ModuleList(
            [TransformerLayer(d_model, n_heads, dropout) for _ in range(3)]
        )

    # Single pass (no CFG)
    def _forward_once(
        self, x: Tensor, time: Tensor, embedding: Optional[Tensor]
    ) -> Tensor:
        """
        x: (B, 1, L) or (B, L)
        time: (B,) or (B,1)
        embedding: (B, T, cond_dim) or (B, cond_dim) or None
        returns: (B, 1, L)
        """
        if x.dim() == 2:
            x = x.unsqueeze(1)  # (B, 1, L)
        B, C, L = x.shape
        assert C == self.in_channels, "x must be [B, 1, L]"
        assert L == self.seq_len, "sequence length mismatch"

        device = x.device

        # Tokens along time
        x_tok = x.transpose(1, 2)  # (B, L, 1) as tokens
        x_tok = self.in_proj(x_tok)  # (B, L, d)

        # Add positional encodings and time bias
        x_tok = x_tok + sinusoidal_positional_encoding(
            L, self.d_model, device
        )
        t_emb = self.time_pos_emb(time.float())
        t_vec = self.time_mlp(t_emb)  # (B, d)
        x_tok = x_tok + t_vec.unsqueeze(1)  # broadcast over L

        # Prepare cond memory (optional)
        cond_mem: Optional[Tensor] = None
        if embedding is not None:
            if embedding.dim() == 2:
                # (B, D) -> (B, 1, D)
                embedding = embedding.unsqueeze(1)
            # (B, T, cond_dim) -> proj -> add PE
            cond_mem = self.cond_proj(embedding)  # (B, T, d)
            cond_mem = cond_mem + sinusoidal_positional_encoding(
                cond_mem.size(1), self.d_model, device
            )

        # Pass through 3 transformer layers
        for layer in self.layers:
            x_tok = layer(x_tok, cond_mem)

        # Back to (B, C, L)
        out = self.out_proj(x_tok).transpose(1, 2)  # (B, C_out, L)
        return out  # keep (B, 1, L)

    def forward(
        self,
        x: Tensor,
        time: Tensor,
        embedding_mask_proba: float = 0.1,
        embedding: Optional[Tensor] = None,
        embedding_scale: float = 1.0,
    ) -> Tensor:
        """
        Keep UNet-like signature, but this is just a 3-layer transformer.

        x: (B, 1, 64) or (B, 64)
        time: (B,) or (B,1)
        embedding: (B, T, 128) or (B, 128) or None
        embedding_mask_proba: used only during training (CFG masking)
        embedding_scale: classifier-free guidance scale (eval-time)
        returns: (B, 1, 64)
        """
        B = x.shape[0]
        device = x.device

        # Classifier-free guidance masking (train-time)
        if (
            self.training
            and embedding is not None
            and embedding_mask_proba > 0.0
        ):
            mask = torch.rand(B, device=device) < embedding_mask_proba
            if mask.any():
                embedding = embedding.clone()
                if embedding.dim() == 2:
                    # (B, D)
                    embedding[mask] = self.null_cond
                elif embedding.dim() == 3:
                    # (B, T, D)
                    T = embedding.size(1)
                    null_seq = (
                        self.null_cond.view(1, 1, -1)
                        .expand(int(mask.sum().item()), T, -1)
                        .to(device)
                    )
                    embedding[mask] = null_seq
                else:
                    raise ValueError("Unsupported embedding dim for masking")

        # CFG (eval-time two passes)
        if (not self.training) and embedding is not None and (
            embedding_scale != 1.0
        ):
            out_cond = self._forward_once(x, time, embedding)
            # build null embedding matching shape
            if embedding.dim() == 2:
                null_emb = self.null_cond.unsqueeze(0).expand(B, -1)
            elif embedding.dim() == 3:
                T = embedding.size(1)
                null_emb = (
                    self.null_cond.unsqueeze(0)
                    .unsqueeze(1)
                    .expand(B, T, -1)
                )
            else:
                raise ValueError("Unsupported embedding dim for CFG")
            out_null = self._forward_once(x, time, null_emb)
            return out_null + (out_cond - out_null) * embedding_scale

        # Single pass
        return self._forward_once(x, time, embedding)


# # Quick check
# if __name__ == "__main__":
#     torch.manual_seed(0)
#     model = UNet1DTransformerSimple(
#         seq_len=64, in_channels=1, out_channels=1, d_model=256, n_heads=8
#     )
#     x = torch.randn(2, 1, 64)
#     t = torch.randint(0, 1000, (2,))
#     cond = torch.randn(2, 20, 128)
#     y = model(x, t, embedding=cond, embedding_scale=1.2)
#     print(y.shape)  # torch.Size([2, 1, 64])

	import math
	from typing import Optional

	import torch
	import torch.nn as nn
	from torch import Tensor


	# -----------------------
	# Positional encodings
	# -----------------------
	def sinusoidal_positional_encoding(
	length: int, dim: int, device: torch.device
	) -> Tensor:
	assert dim % 2 == 0, "positional dim must be even"
	pos = torch.arange(length, device=device).float().unsqueeze(1) # (L, 1)
	half = dim // 2
	div = torch.exp(
	torch.arange(0, half, device=device).float()
	* (-math.log(10000.0) / max(1, half - 1))
	) # (half,)
	pe = torch.cat([torch.sin(pos * div), torch.cos(pos * div)], dim=1)
	return pe.unsqueeze(0) # (1, L, dim)


	class SinusoidalPosEmb(nn.Module):
	def __init__(self, dim: int):
	super().__init__()
	assert dim % 2 == 0
	self.dim = dim

	def forward(self, t: Tensor) -> Tensor:
	if t.dim() == 2 and t.size(1) == 1:
	t = t[:, 0]
	half = self.dim // 2
	device = t.device
	freq = torch.exp(
	-math.log(10000)
	* torch.arange(half, device=device).float()
	/ (half - 1)
	)
	args = t.float().unsqueeze(1) * freq.unsqueeze(0)
	emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
	return emb # (B, dim)


	# -----------------------
	# Transformer building blocks
	# -----------------------
	class TransformerLayer(nn.Module):
	"""
	Minimal layer:
	- self-attn on x
	- cross-attn x -> cond (if provided)
	- MLP
	All with pre-norm and residuals.
	"""

	def __init__(self, d_model: int, n_heads: int = 8, dropout: float = 0.0):
	super().__init__()
	self.norm1 = nn.LayerNorm(d_model)
	self.self_attn = nn.MultiheadAttention(
	embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
	)
	self.drop1 = nn.Dropout(dropout)

	self.norm2 = nn.LayerNorm(d_model)
	self.cross_attn = nn.MultiheadAttention(
	embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
	)
	self.drop2 = nn.Dropout(dropout)

	self.norm3 = nn.LayerNorm(d_model)
	self.mlp = nn.Sequential(
	nn.Linear(d_model, d_model * 4),
	nn.GELU(),
	nn.Dropout(dropout),
	nn.Linear(d_model * 4, d_model),
	nn.Dropout(dropout),
	)

	def forward(self, x: Tensor, cond: Optional[Tensor]) -> Tensor:
	# x: (B, L, d), cond: (B, T, d) or None
	y = self.norm1(x)
	y, _ = self.self_attn(y, y, y, need_weights=False)
	x = x + self.drop1(y)

	if cond is not None:
	z = self.norm2(x)
	z, _ = self.cross_attn(z, cond, cond, need_weights=False)
	x = x + self.drop2(z)

	w = self.norm3(x)
	x = x + self.mlp(w)
	return x


	# -----------------------
	# Just-a-Transformer "UNet1D" (no up/down), 3 layers
	# Keeps UNet-like forward signature incl. CFG knobs.
	# -----------------------
	class UNet1DTransformerSimple(nn.Module):
	def __init__(
	self,
	seq_len: int = 64,
	in_channels: int = 1,
	out_channels: int = 1,
	d_model: int = 256,
	n_heads: int = 8,
	dropout: float = 0.0,
	cond_dim: int = 128, # embedding is [B, T, 128]
	time_emb_dim: int = 128,
	):
	super().__init__()
	self.seq_len = seq_len
	self.in_channels = in_channels
	self.out_channels = out_channels
	self.d_model = d_model
	self.cond_dim = cond_dim

	# Input/output projections
	self.in_proj = nn.Linear(in_channels, d_model)
	self.out_proj = nn.Linear(d_model, out_channels)

	# Time embedding -> d_model bias
	self.time_pos_emb = SinusoidalPosEmb(time_emb_dim)
	self.time_mlp = nn.Sequential(
	nn.Linear(time_emb_dim, d_model * 4),
	nn.SiLU(),
	nn.Linear(d_model * 4, d_model),
	)

	# Cond projection to d_model (for cross-attention memory)
	self.cond_proj = nn.Linear(cond_dim, d_model)
	# Learnable null cond for classifier-free guidance
	self.null_cond = nn.Parameter(torch.zeros(cond_dim))

	# Exactly 3 transformer layers
	self.layers = nn.ModuleList(
	[TransformerLayer(d_model, n_heads, dropout) for _ in range(3)]
	)

	# Single pass (no CFG)
	def _forward_once(
	self, x: Tensor, time: Tensor, embedding: Optional[Tensor]
	) -> Tensor:
	"""
	x: (B, 1, L) or (B, L)
	time: (B,) or (B,1)
	embedding: (B, T, cond_dim) or (B, cond_dim) or None
	returns: (B, 1, L)
	"""
	if x.dim() == 2:
	x = x.unsqueeze(1) # (B, 1, L)
	B, C, L = x.shape
	assert C == self.in_channels, "x must be [B, 1, L]"
	assert L == self.seq_len, "sequence length mismatch"

	device = x.device

	# Tokens along time
	x_tok = x.transpose(1, 2) # (B, L, 1) as tokens
	x_tok = self.in_proj(x_tok) # (B, L, d)

	# Add positional encodings and time bias
	x_tok = x_tok + sinusoidal_positional_encoding(
	L, self.d_model, device
	)
	t_emb = self.time_pos_emb(time.float())
	t_vec = self.time_mlp(t_emb) # (B, d)
	x_tok = x_tok + t_vec.unsqueeze(1) # broadcast over L

	# Prepare cond memory (optional)
	cond_mem: Optional[Tensor] = None
	if embedding is not None:
	if embedding.dim() == 2:
	# (B, D) -> (B, 1, D)
	embedding = embedding.unsqueeze(1)
	# (B, T, cond_dim) -> proj -> add PE
	cond_mem = self.cond_proj(embedding) # (B, T, d)
	cond_mem = cond_mem + sinusoidal_positional_encoding(
	cond_mem.size(1), self.d_model, device
	)

	# Pass through 3 transformer layers
	for layer in self.layers:
	x_tok = layer(x_tok, cond_mem)

	# Back to (B, C, L)
	out = self.out_proj(x_tok).transpose(1, 2) # (B, C_out, L)
	return out # keep (B, 1, L)

	def forward(
	self,
	x: Tensor,
	time: Tensor,
	embedding_mask_proba: float = 0.1,
	embedding: Optional[Tensor] = None,
	embedding_scale: float = 1.0,
	) -> Tensor:
	"""
	Keep UNet-like signature, but this is just a 3-layer transformer.

	x: (B, 1, 64) or (B, 64)
	time: (B,) or (B,1)
	embedding: (B, T, 128) or (B, 128) or None
	embedding_mask_proba: used only during training (CFG masking)
	embedding_scale: classifier-free guidance scale (eval-time)
	returns: (B, 1, 64)
	"""
	B = x.shape[0]
	device = x.device

	# Classifier-free guidance masking (train-time)
	if (
	self.training
	and embedding is not None
	and embedding_mask_proba > 0.0
	):
	mask = torch.rand(B, device=device) < embedding_mask_proba
	if mask.any():
	embedding = embedding.clone()
	if embedding.dim() == 2:
	# (B, D)
	embedding[mask] = self.null_cond
	elif embedding.dim() == 3:
	# (B, T, D)
	T = embedding.size(1)
	null_seq = (
	self.null_cond.view(1, 1, -1)
	.expand(int(mask.sum().item()), T, -1)
	.to(device)
	)
	embedding[mask] = null_seq
	else:
	raise ValueError("Unsupported embedding dim for masking")

	# CFG (eval-time two passes)
	if (not self.training) and embedding is not None and (
	embedding_scale != 1.0
	):
	out_cond = self._forward_once(x, time, embedding)
	# build null embedding matching shape
	if embedding.dim() == 2:
	null_emb = self.null_cond.unsqueeze(0).expand(B, -1)
	elif embedding.dim() == 3:
	T = embedding.size(1)
	null_emb = (
	self.null_cond.unsqueeze(0)
	.unsqueeze(1)
	.expand(B, T, -1)
	)
	else:
	raise ValueError("Unsupported embedding dim for CFG")
	out_null = self._forward_once(x, time, null_emb)
	return out_null + (out_cond - out_null) * embedding_scale

	# Single pass
	return self._forward_once(x, time, embedding)


	# # Quick check
	# if __name__ == "__main__":
	# torch.manual_seed(0)
	# model = UNet1DTransformerSimple(
	# seq_len=64, in_channels=1, out_channels=1, d_model=256, n_heads=8
	# )
	# x = torch.randn(2, 1, 64)
	# t = torch.randint(0, 1000, (2,))
	# cond = torch.randn(2, 20, 128)
	# y = model(x, t, embedding=cond, embedding_scale=1.2)
	# print(y.shape) # torch.Size([2, 1, 64])