| 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):
|
| if self.inv_freq_h is None or self.inv_freq_h.device.type == "meta" or (self.inv_freq_h == 0).all():
|
| self.inv_freq_h = 1.0 / (self.base ** (torch.arange(0, self.dim_h, 2, dtype=torch.float32, device=device) / self.dim_h))
|
| self.inv_freq_w = 1.0 / (self.base ** (torch.arange(0, self.dim_w, 2, dtype=torch.float32, device=device) / self.dim_w))
|
|
|
| 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:
|
| if not hasattr(self, "hadamard_mat") or self.hadamard_mat is None or self.hadamard_mat.device.type == "meta" or (self.hadamard_mat == 0).all() or self.hadamard_mat.abs().max() > 10.0:
|
| hadamard_mat = get_hadamard_matrix(self.dim, dtype=torch.float32).to(x.device)
|
| self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)
|
| 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)
|
| elif os.path.exists(bin_file):
|
| state_dict = torch.load(bin_file, map_location="cpu")
|
| else:
|
| model = super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
|
| for module in model.modules():
|
| if type(module).__name__ == "RotaryEmbedding2D":
|
| module.inv_freq_h = 1.0 / (module.base ** (torch.arange(0, module.dim_h, 2, dtype=torch.float32) / module.dim_h))
|
| module.inv_freq_w = 1.0 / (module.base ** (torch.arange(0, module.dim_w, 2, dtype=torch.float32) / module.dim_w))
|
| elif type(module).__name__ == "HadamardMLP":
|
| module.register_buffer("hadamard_mat", get_hadamard_matrix(module.dim, dtype=torch.float32), persistent=False)
|
| return model
|
|
|
| model_keys = set(model.state_dict().keys())
|
| st_keys = set(state_dict.keys())
|
| if not model_keys.intersection(st_keys):
|
| if any(k.startswith("model.") for k in st_keys):
|
| state_dict = {k[6:] if k.startswith("model.") else k: v for k, v in state_dict.items()}
|
| elif any(k.startswith("model.") for k in model_keys):
|
| state_dict = {f"model.{k}": v for k, v in state_dict.items()}
|
|
|
| model.load_state_dict(state_dict, strict=True)
|
|
|
| for module in model.modules():
|
| if type(module).__name__ == "RotaryEmbedding2D":
|
| module.inv_freq_h = 1.0 / (module.base ** (torch.arange(0, module.dim_h, 2, dtype=torch.float32) / module.dim_h))
|
| module.inv_freq_w = 1.0 / (module.base ** (torch.arange(0, module.dim_w, 2, dtype=torch.float32) / module.dim_w))
|
| elif type(module).__name__ == "HadamardMLP":
|
| module.register_buffer("hadamard_mat", get_hadamard_matrix(module.dim, dtype=torch.float32), persistent=False)
|
|
|
| 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)
|
|
|
