noisy = noise * (1 - t) + pixel_values * t
v_pred = model.forward(noisy, t, ctx)
v_target = pixel_values - noise
loss = torch.nn.functional.mse_loss(v_pred, v_target)

Inference

python

@torch.no_grad()
def inference(model: DiT, device=None, steps=50):
    tokenizer = AutoTokenizer.from_pretrained('nebulette/booru-character-aware-tokenizer')
    ctx = torch.tensor(tokenizer.encode('portrait')).unsqueeze(0).to(device)
    xt = torch.randn((1, 3, 48, 48), device=device)

    # Generate time steps from 0 to 1.
    time_steps = torch.linspace(0.0, 1.0, steps + 1, device=device)

    # Iterate through time steps.
    for t in time_steps:
        t = t.unsqueeze(0)
        # Predict the velocity at point (x_t, t) using the model.
        v_pred = model.forward(xt, t, ctx)

        # Update the state based on the predicted velocity.
        xt = xt + v_pred * (1 / steps)

    # Convert CIELAB → RGB.
    lab = torch.clamp(xt[0], -1, 1).cpu().numpy()
    L = (lab[0] + 1) * 50
    a = lab[1] * 128
    b = lab[2] * 128
    rgb = color.lab2rgb(np.stack([L, a, b], axis=-1)) * 255.0

    return Image.fromarray(rgb.astype(np.uint8))

References

3588028.3603685
2503.16397
2508.02324 (Figure 8)

nebulette/segmentation-aware-jit-b

JiT-B

What's new

Future

Training

Inference

References

nebulette/
segmentation-aware-jit-b