class XPredictionModel(nn.Module):
"""
Diffusion model with x-prediction head.
Can be Transformer, UNet, etc.; example: Transformer-based.
"""
def __init__(self, input_channels=3, patch_size=4, hidden_dim=768, num_layers=12):
super().__init__()
self.patch_size = patch_size
self.hidden_dim = hidden_dim
self.patch_embed = nn.Linear(3 * patch_size * patch_size, hidden_dim)
self.time_embed = nn.Sequential(
nn.Linear(1, hidden_dim // 2),
nn.GELU(),
nn.Linear(hidden_dim // 2, hidden_dim)
)
self.transformer = nn.TransformerEncoder(
nn.TransformerEncoderLayer(hidden_dim, nhead=8, dim_feedforward=3072),
num_layers=num_layers
)
self.x_pred_head = nn.Linear(hidden_dim, 3 * patch_size * patch_size)
def forward(self, x_t, t):
"""
x_t: noisy input (batch, 3, H, W)
t: timestep (batch,)
Returns: predicted clean image (batch, 3, H, W)
"""
batch, channels, height, width = x_t.shape
patches = self._patchify(x_t)
patch_embeds = self.patch_embed(patches)
time_emb = self.time_embed(t.unsqueeze(-1).float())
time_emb = time_emb.unsqueeze(1)
patch_embeds = patch_embeds + time_emb
hidden = self.transformer(patch_embeds)
clean_patches = self.x_pred_head(hidden)
x_0_pred = self._unpatchify(clean_patches, height, width)
return x_0_pred
def _patchify(self, x):
"""Convert image to non-overlapping patches."""
batch, channels, height, width = x.shape
x = x.reshape(
batch,
channels,
height // self.patch_size,
self.patch_size,
width // self.patch_size,
self.patch_size
)
x = x.permute(0, 2, 4, 1, 3, 5).contiguous()
x = x.reshape(batch, -1, channels * self.patch_size * self.patch_size)
return x
def _unpatchify(self, patches, height, width):
"""Reconstruct image from patches."""
batch, num_patches, patch_dim = patches.shape
channels = 3
patches = patches.reshape(
batch,
height // self.patch_size,
width // self.patch_size,
channels,
self.patch_size,
self.patch_size
)
patches = patches.permute(0, 3, 1, 4, 2, 5).contiguous()
x = patches.reshape(batch, channels, height, width)
return x
def train_x_prediction_diffusion(model, dataloader, num_epochs=100, lr=1e-4):
"""
Train diffusion model with x-prediction.
"""
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
diffusion = XPredictionDiffusion(timesteps=1000)
for epoch in range(num_epochs):
total_loss = 0
for x_0 in dataloader:
batch_size = x_0.shape[0]
t = torch.randint(0, diffusion.timesteps, (batch_size,))
x_t, noise = diffusion.add_noise(x_0, t)
x_0_pred = model(x_t, t)
loss = diffusion.loss_x_prediction(x_0, t, x_0_pred)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch}: avg_loss={avg_loss:.4f}")
@torch.no_grad()
def generate_x_prediction(model, diffusion, shape, num_steps=100):
"""
Generate samples via iterative x-prediction denoising.
"""
x_t = torch.randn(shape)
timesteps = torch.linspace(999, 0, num_steps).long()
for t in timesteps:
t_batch = torch.full((x_t.shape[0],), t)
x_0_pred = model(x_t, t_batch)
alpha_t = diffusion.alphas_cumprod[t]
alpha_prev = diffusion.alphas_cumprod[t - 1] if t > 0 else 1.0
c1 = torch.sqrt((1 - alpha_prev) / (1 - alpha_t))
c2 = torch.sqrt(alpha_prev / alpha_t)
x_t = c2 * x_0_pred + c1 * (x_t - torch.sqrt(1 - alpha_t) * x_0_pred)
if t > 0:
noise = torch.randn_like(x_t)
sigma_t = torch.sqrt((1 - alpha_prev) / (1 - alpha_t) * (1 - alpha_t / alpha_prev))
x_t = x_t + sigma_t * noise
return x_t.clamp(-1, 1)