菜单
Train a unified autoregressive model to generate images and text by directly handling continuous image tokens with flow matching, eliminating the need for quantization or separate diffusion models.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill nextstep-1-autoregressive-images命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
基于 SOC 职业分类
正在显示 SKILL.md
| name | nextstep-1-autoregressive-images |
| title | NextStep-1: Autoregressive Image Generation with Continuous Tokens |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2508.10711 |
| keywords | ["image-generation","autoregressive-modeling","continuous-tokens","flow-matching","text-to-image"] |
| description | Train a unified autoregressive model to generate images and text by directly handling continuous image tokens with flow matching, eliminating the need for quantization or separate diffusion models. |
Most text-to-image models use separate pipelines: text encoders, diffusion models, and discrete image tokenizers. NextStep-1 unifies image and text generation in a single autoregressive architecture that directly processes continuous image tokens using flow matching instead of diffusion.
The key innovation is combining a large autoregressive language model (14B parameters) with a smaller flow matching head (157M parameters) to predict both discrete text tokens and continuous image tokens in a unified next-token prediction framework.
Convert images into continuous token representations using a learned encoder. These are embedded directly into the autoregressive sequence without quantization.
import torch
import torch.nn as nn
from transformers import AutoModel
class ContinuousImageTokenizer(nn.Module):
"""
Encodes images to continuous token embeddings
Similar to VQ-VAE but without quantization
"""
def __init__(self, image_size=256, token_dim=768, num_tokens=1024):
super().__init__()
self.image_size = image_size
self.token_dim = token_dim
self.num_tokens = num_tokens
.encoder = nn.Sequential(
nn.Conv2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.Conv2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.Conv2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
)
.to_tokens = nn.Linear( * * , num_tokens * token_dim)
():
features = .encoder(images)
batch_size = features.shape[]
features = features.view(batch_size, -)
tokens = .to_tokens(features)
tokens = tokens.view(batch_size, .num_tokens, .token_dim)
tokens
():
batch_size = tokens.shape[]
flat = tokens.view(batch_size, -)
features = .to_tokens.weight.T @ flat.T
features = features.view(batch_size, , , )
decoder = nn.Sequential(
nn.ConvTranspose2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.ConvTranspose2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.ConvTranspose2d(, , kernel_size=, stride=, padding=),
nn.Tanh(),
)
images = decoder(features)
images
Combine text tokens and continuous image tokens into single sequences that the autoregressive model can process.
def create_mixed_sequence(text_ids, image_tokens, tokenizer, image_token_start=50257):
"""
Create interleaved sequence of text and image tokens.
Format: [text_tokens] [IMAGE_START] [image_tokens] [IMAGE_END]
Args:
text_ids: [seq_len] text token IDs
image_tokens: [num_image_tokens, token_dim] continuous embeddings
tokenizer: text tokenizer
image_token_start: special token ID marking image start
Returns:
input_ids: [total_seq_len] mixed token sequence (text IDs as scalars)
embeddings: [total_seq_len, embedding_dim] actual embeddings
"""
text_embedding_dim = 768
image_token_dim = 768
# Text embeddings
text_embeddings = tokenizer.transformer.wte(text_ids) # [seq_len, 768]
# Insert image tokens
input_ids = torch.cat([
text_ids,
torch.tensor([image_token_start], device=text_ids.device),
torch.arange(50257 + 1, 50257 + image_tokens.shape[0] + 1,
device=text_ids.device),
torch.tensor([image_token_start + 1], device=text_ids.device)
])
# Combine embeddings
image_start_embed = text_embeddings[-1:] * 0.5 # Learned marker
embeddings = torch.cat([
text_embeddings,
image_start_embed,
image_tokens, # ← Continuous embeddings inserted directly
image_start_embed
], dim=0)
return input_ids, embeddings
Build a transformer that predicts next tokens, treating continuous image tokens like any other embedding.
class UnifiedAutoregressive(nn.Module):
"""
14B parameter autoregressive model for text + continuous image tokens
"""
def __init__(self, vocab_size=50257, hidden_size=1024, num_layers=24,
num_image_tokens=1024, image_token_dim=768):
super().__init__()
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_image_tokens = num_image_tokens
# Text embedding
self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
# Transformer decoder
self.transformer = nn.ModuleList([
nn.TransformerDecoderLayer(
d_model=hidden_size, nhead=16,
dim_feedforward=4096, batch_first=True
) for _ in range(num_layers)
])
# Output heads for next token prediction
self.text_head = nn.Linear(hidden_size, vocab_size) # Discrete text tokens
self.image_token_head = nn.Linear(hidden_size, image_token_dim) # Continuous image
def forward(self, input_ids, embeddings):
"""
Args:
input_ids: [batch, seq_len] (for masking/identifying token types)
embeddings: [batch, seq_len, hidden_size] mixed embeddings
Returns:
text_logits: [batch, seq_len, vocab_size] for text token prediction
image_logits: [batch, seq_len, image_token_dim] for image token prediction
"""
hidden = embeddings
# Forward through transformer layers
for layer in self.transformer:
hidden = layer(hidden, hidden)
# Predict next tokens
text_logits = self.text_head(hidden)
image_logits = self.image_token_head(hidden)
return text_logits, image_logits
Use flow matching instead of diffusion to train the continuous image token predictions.
class FlowMatchingHead(nn.Module):
"""
Flow matching objective for continuous image tokens.
Simpler than diffusion, more stable than MSE loss.
"""
def __init__(self, token_dim=768, hidden_size=1024):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(token_dim + hidden_size, 512),
nn.GELU(),
nn.Linear(512, token_dim)
)
def forward(self, target_tokens, context_hidden, t=None):
"""
Flow matching loss: learn to map noise to data
target: continuous image tokens [batch, num_tokens, token_dim]
context: model hidden states [batch, seq_len, hidden_size]
"""
batch_size, num_tokens, token_dim = target_tokens.shape
# Flow matching: learn velocity field from noise to data
noise = torch.randn_like(target_tokens)
t = torch.rand(batch_size, 1, 1, device=target_tokens.device)
# Interpolate: z_t = (1-t) * noise + t * target
z_t = (1 - t) * noise + t * target_tokens
# Predict velocity
context_repeated = context_hidden[:, -num_tokens:, :].unsqueeze(1)
velocity_pred = self.mlp(torch.cat([z_t, context_repeated.expand_as(z_t)], dim=-1))
# Flow matching loss: velocity should match (target - noise)
target_velocity = target_tokens - noise
flow_loss = ((velocity_pred - target_velocity) ** 2).mean()
return flow_loss
Combine text prediction loss (cross-entropy) with flow matching loss for image tokens.
def train_nextstep(model, flow_head, image_tokenizer, dataloader, num_epochs=10):
"""
Train unified autoregressive model with mixed objectives
"""
optimizer = torch.optim.AdamW(
list(model.parameters()) + list(flow_head.parameters()) +
list(image_tokenizer.parameters()),
lr=1e-4
)
for epoch in range(num_epochs):
for batch_idx, (images, captions) in enumerate(dataloader):
# Encode images to continuous tokens
image_tokens = image_tokenizer.encode(images) # [batch, num_tokens, 768]
# Tokenize text
text_ids = tokenizer.encode(captions)
# Create mixed sequences
input_ids, embeddings = create_mixed_sequence(
text_ids, image_tokens, tokenizer
)
# Forward pass
text_logits, image_logits = model(input_ids, embeddings)
# Text prediction loss (standard cross-entropy)
text_loss = F.cross_entropy(
text_logits.view(-1, model.vocab_size),
text_ids.view(-1)
)
# Image token prediction with flow matching
# Extract image token portion of context
image_context = embeddings[:, :text_ids.shape[1], :]
flow_loss = flow_head(image_tokens, image_context)
# Combined loss
total_loss = text_loss + 0.5 * flow_loss
optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
if batch_idx % 100 == 0:
print(f"Epoch {epoch}, Step {batch_idx}")
print(f" Text Loss: {text_loss:.4f}, Flow Loss: {flow_loss:.4f}")
Use the trained model to generate images autoregressively from text prompts.
@torch.no_grad()
def generate_image(model, flow_head, image_tokenizer, prompt, max_new_tokens=1024):
"""
Generate image autoregressively from text prompt
"""
# Encode prompt
input_ids = tokenizer.encode(prompt)
# Add image start token
input_ids = torch.cat([
input_ids,
torch.tensor([50257], device=input_ids.device) # IMAGE_START
])
# Autoregressive generation of image tokens
generated_tokens = []
hidden_state = None
for step in range(max_new_tokens):
embeddings = model.embed_tokens(input_ids.unsqueeze(0))
text_logits, image_logits = model(input_ids.unsqueeze(0), embeddings)
# Sample next image token
next_token_logits = image_logits[0, -1, :]
next_token = torch.randn_like(next_token_logits) * 0.1 + next_token_logits
generated_tokens.append(next_token)
if step % 100 == 0:
print(f"Generated {step}/{max_new_tokens} image tokens")
# Stack generated image tokens
image_tokens = torch.stack(generated_tokens) # [num_tokens, 768]
# Decode to image
image = image_tokenizer.decode(image_tokens.unsqueeze(0))
return image
NextStep-1 (2508.10711): https://arxiv.org/abs/2508.10711
Unified autoregressive model for text and continuous image tokens with flow matching, eliminating quantization loss and diffusion overhead while enabling efficient image generation and editing.