Codon Motif Vision 1

模型介绍

Codon Motif Vision 1 是一个实验性的视觉 Tokenizer (Visual Tokenizer) 模型，可以支持任意尺度图片。

技术规格

该模型基于 VQ-GAN 架构，并融合了以下关键技术：

• 量化 (Quantizer): 采用 LFQ (Lookup Free Quantization)，一种无查表的向量量化技术，计算包含熵 (Entropy) 和困惑度 (Perplexity)。
• 位置编码: 集成 2D RoPE (Rotary Positional Embedding) 二维旋转位置编码，增强空间感知。
• 注意力机制: Encoder 和 Decoder 均使用 Spatial Multi-Head Attention (空间多头注意力)。
• 编码器 (Encoder): 结合 ConvBlock 和 ResBasicBlock，支持动态输入尺寸。
• 解码器 (Decoder): 使用 PixelShuffle 进行上采样重建。

模型参数与文件

文件清单

默认配置

以下为模型初始化默认参数（即 v1 标准配置）：

• 输入/输出通道: 3 (RGB)
• Patch Size: 16
• Latent Dim: 256
• Codebook Dim: 16
• Encoder: Hidden Dim 256, 1 ResBlock, 4 Heads
• Decoder: Hidden Dim 256, 3 ResBlocks, 4 Heads
• RoPE Max Len: 4096

使用方法

1. 安装依赖

确保已安装 orbit-torch 库：

pip install orbit-torch

2. 模型加载

使用 load_pretrained 方法加载权重。权重分为全量权重和部分权重（Encoder, Decoder, Quantizer）。

导入模块

from orbit.model.motif.vision.v1 import MotifV1

场景 A：加载全量模型

model = MotifV1()
model.load_pretrained('motif-v1.safetensors')

场景 B：仅使用 Encoder

注意：在使用 Encoder 提取 Token 时，必须同时加载 Quantizer。通过 MotifV1 实例访问子模块进行加载。

# 1. 实例化主模型
model = MotifV1()

# 2. 分别加载权重
model.encoder.load_pretrained('motif-v1_encoder.safetensors')
model.quantizer.load_pretrained('motif-v1_quantizer.safetensors')

# 3. 使用示例 (使用封装好的 encode 方法)
# x = torch.randn(1, 3, 256, 256) # [B, 3, H, W]
# indices, mask, z_q = model.encode(x)
# print(indices.shape) # [B, H, W]

场景 C：仅使用 Decoder

注意：在使用 Decoder 重建图像时，必须同时加载 Quantizer（用于将索引还原为向量）。

# 1. 实例化主模型
model = MotifV1()

# 2. 分别加载权重
model.decoder.load_pretrained('motif-v1_decoder.safetensors')
model.quantizer.load_pretrained('motif-v1_quantizer.safetensors')

# 3. 使用示例 (使用封装好的 decode 方法)
# indices = ... # [B, H, W]
# reconstruction = model.decode(indices)