下面给出一个相对完整、可运行的示例性微调脚本，演示从原始 JSONL 数据到生成最终体素结果的完整流程。需要注意的是，以下示例代码主要用于教学和参考，实际项目中仍需根据你自身的数据格式、超参数需求和硬件情况进行调整。

目录

1. 依赖安装
2. 数据准备
3. 模型定义
4. 训练循环
5. 采样与导出
6. 完整示例脚本

1. 依赖安装

在使用此脚本前，请确保已经安装相关依赖：

bash
pip install shap-e
pip install transformers
pip install torch torchvision

• shap-e：OpenAI 提供的 Shap-E 库
• transformers：Hugging Face 的 Transformer 工具包
• torch / torchvision：PyTorch

2. 数据准备

假设你有一个包含 Minecraft 体素描述和坐标的 JSONL 文件（例如 training_data/10001-1.jsonl），格式示例如下：

json
{"input": "一个红色的羊毛方块", "output": "red_wool 10 10 10"}
{"input": "一个石头搭建的房子", "output": "stone 0 0 0\nstone 1 0 0\n..."}
...

其中：

• "input" 为自然语言描述
• "output" 是若干行描述体素坐标信息，每行格式形如：方块名称 x y z

你还需要一个 blockids.txt 文件，里面定义了方块名称与数字 ID 的对应关系。例如：

python
block_ids = {
    0: "air",
    1: "stone",
    2: "dirt",
    3: "red_wool",
    ...
}

注意：上面只是一个示例，你需要根据实际的 Minecraft 方块种类自行定义映射。

然后我们将把这些 JSONL 数据转换成两个 PyTorch 文件：

1. descriptions.pt：保存的是原始文本字符串列表（或者批量）
2. voxels.pt：保存的是体素稀疏表示（或构造成密集 3D 张量后再保存）

下面给出一个示例脚本，将 JSONL 转为可用的训练数据（稀疏体素列表）并保存下来。你可以单独保存成 prepare_data.py 或者直接放在你的主脚本里运行也行。

python
import os
import json
import torch

from transformers import CLIPTokenizer

# ========== 1. 定义路径 ==========
training_data_path = 'training_data/10001-1.jsonl'  # 你的 JSONL 文件
block_ids_path = 'blockids.txt'                     # 存放 block_ids 映射的 Python 脚本
output_folder = 'processed_output'                  # 输出文件夹

os.makedirs(output_folder, exist_ok=True)

# ========== 2. 加载 block_ids ==========
block_ids = {}
with open(block_ids_path, 'r', encoding='utf-8') as f:
    exec(f.read())  # 这里会执行 block_ids = { ... }

# ========== 3. 初始化 tokenizer ==========
# 如果后续还要用 CLIPTokenizer 做分词，这里可以提前加载
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")

# ========== 4. 准备容器 ==========
description_list = []  # 用于保存原始文本
voxel_list = []        # 用于保存[(block_id, x, y, z), (block_id, x, y, z), ...]

# ========== 5. 读取 JSONL 文件 ==========
with open(training_data_path, 'r', encoding='utf-8') as f:
    for line in f:
        data = json.loads(line.strip())
        description = data.get('input', '')  # 读取自然语言描述
        voxel_data = data.get('output', '')  # 读取体素坐标
        
        # 将描述直接存储为字符串（后面在模型训练阶段再进行分词）
        description_list.append(description)
        
        # 处理体素数据，每行格式: block_name x y z
        voxel_lines = voxel_data.strip().split('\n')
        voxel_ids = []
        for voxel_line in voxel_lines:
            parts = voxel_line.strip().split()
            if len(parts) == 4:
                block_name, x, y, z = parts
                # 根据 block_ids 找到对应的数字 ID
                block_id = None
                for k, v in block_ids.items():
                    if v == block_name:
                        block_id = k
                        break
                if block_id is None:
                    print(f"[警告] 未找到方块名称 '{block_name}' 的映射，跳过。")
                    continue
                voxel_ids.append([int(block_id), int(x), int(y), int(z)])
            else:
                print(f"[警告] 无效的体素行: {voxel_line}")
        # voxel_ids 是一个 (N, 4) 的列表
        voxel_list.append(voxel_ids)

# ========== 6. 保存数据 ==========
torch.save(description_list, os.path.join(output_folder, 'descriptions.pt'))
torch.save(voxel_list, os.path.join(output_folder, 'voxels.pt'))

print("数据预处理完成，已生成 descriptions.pt 和 voxels.pt。")

在这里，我们将描述保存在一个 Python 列表（每个元素是一个字符串），体素信息保存在一个 Python 列表（每个元素又是 (N,4) 的坐标与 block_id 的稀疏形式）。

3. 模型定义

下面是简化的模型定义示例，包括：

1. 文本编码器：MinecraftCLIPTextEncoder
2. Voxel 编码器：VoxelEncoder（将体素数据编码到潜变量）
3. Voxel 解码器：VoxelDecoder（将潜变量解码回体素分布）
4. 扩散模型：借助 shap_e.diffusion.gaussian_diffusion 中的 diffusion_from_config

稠密体素/稀疏体素转换

在训练 VoxelEncoder 时，如果需要一个固定分辨率（如 64x64x64）的 3D 网格，就需要把 (N,4) 的稀疏表示转换为一个 dense_tensor[64,64,64]，其中每个体素存储其 block_id。

• 如果超出体素边界会被忽略
• EMPTY_BLOCK_ID = 0 表示“空气”或“无方块”

示例：

python
def voxel_list_to_dense(voxel_ids, size=64, empty_id=0):
    """
    将体素的稀疏表示(一堆 [block_id, x, y, z]) 转为 [size, size, size] 的密集张量。
    """
    volume = torch.full((size, size, size), empty_id, dtype=torch.long)
    for (block_id, vx, vy, vz) in voxel_ids:
        if 0 <= vx < size and 0 <= vy < size and 0 <= vz < size:
            volume[vx, vy, vz] = block_id
    return volume

代码示例

python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from transformers import CLIPTokenizer, CLIPTextModel

from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
from shap_e.models.download import load_config
from shap_e.models.generation.perceiver import SimplePerceiver
from shap_e.models.generation.transformer import Transformer

# ========== 超参数定义 ==========
VOXEL_SIZE = 64        # 体素分辨率
EMPTY_BLOCK_ID = 0     # 空气或空方块
NUM_BLOCK_TYPES = 600  # 假设你有 600 种不同方块，需自行根据 block_ids 配置

# ========== 1. 文本编码器 (CLIP) ==========
class MinecraftCLIPTextEncoder(nn.Module):
    """
    使用 CLIP 文本模型作为底座，再接一个线性层，将输出映射到自定义维度。
    如果你只想要 CLIP 的 pooler_output，可将 proj 层去掉或仅用于扩散条件。
    """
    def __init__(self, clip_model_name="openai/clip-vit-large-patch14", out_dim=512):
        super().__init__()
        self.clip = CLIPTextModel.from_pretrained(clip_model_name)
        self.tokenizer = CLIPTokenizer.from_pretrained(clip_model_name)
        self.proj = nn.Linear(self.clip.config.hidden_size, out_dim)

    def forward(self, text_list):
        """
        text_list: list of strings (原始句子)
        """
        # 将文本转为 token
        inputs = self.tokenizer(text_list, return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(self.clip.device) for k, v in inputs.items()}
        
        outputs = self.clip(**inputs)
        # pooler_output 通常是 [batch_size, hidden_size]
        pooled = outputs.pooler_output  # 或者 outputs.last_hidden_state 的某个池化
        projected = self.proj(pooled)   # [batch_size, out_dim]
        return projected

# ========== 2. 体素编码器 ==========
class VoxelEncoder(nn.Module):
    """
    将体素网格 (batch, 1, D, H, W) -> (batch, d_latent)
    可以根据需要修改网络结构、通道数等。
    """
    def __init__(self, d_latent=512):
        super().__init__()
        self.conv1 = nn.Conv3d(1, 32, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.Conv3d(32, 64, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.Conv3d(64, 128, kernel_size=4, stride=2, padding=1)
        self.fc = nn.Linear(128 * 8 * 8 * 8, d_latent)

    def forward(self, x):
        # x: [batch_size, 1, 64, 64, 64]
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        return x

# ========== 3. 体素解码器 ==========
class VoxelDecoder(nn.Module):
    """
    将潜变量 (batch, d_latent) -> (batch, NUM_BLOCK_TYPES, D, H, W)
    输出 logits，用 argmax 或 softmax 获取每个体素的 block_id。
    """
    def __init__(self, d_latent=512, output_size=64, n_channels=NUM_BLOCK_TYPES):
        super().__init__()
        self.fc = nn.Linear(d_latent, 128 * 8 * 8 * 8)
        self.deconv1 = nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, padding=1)
        self.deconv2 = nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, padding=1)
        self.deconv3 = nn.ConvTranspose3d(32, n_channels, kernel_size=4, stride=2, padding=1)
        self.output_size = output_size

    def forward(self, x):
        x = self.fc(x)  # (batch, 128*8*8*8)
        x = x.view(x.shape[0], 128, 8, 8, 8)  # (batch, 128, 8, 8, 8)
        x = F.relu(self.deconv1(x))          # -> (batch, 64, 16, 16, 16)
        x = F.relu(self.deconv2(x))          # -> (batch, 32, 32, 32, 32)
        x = self.deconv3(x)                  # -> (batch, n_channels, 64, 64, 64)
        return x  # logits

4. 训练循环

1. 读取 descriptions.pt 和 voxels.pt
2. 将 (block_id, x, y, z) 稀疏表示转换为 64x64x64 的 dense 张量
3. 送入 VoxelEncoder，并对照文本嵌入训练扩散模型
4. 每个 epoch 后可以保存模型权重并进行一次采样测试

下面是示例训练循环代码，你可以放到同一个 .py 文件中，也可以拆开多个文件组织。

python
################################
# 载入并构造 DataLoader
################################
class MinecraftDataset(Dataset):
    """
    将 descriptions.pt 与 voxels.pt 拼起来，并在 __getitem__ 中做稀疏->稠密转换。
    """
    def __init__(self, descriptions_path, voxels_path, voxel_size=64):
        super().__init__()
        self.descriptions = torch.load(descriptions_path)  # list of strings
        self.voxels_sparse = torch.load(voxels_path)       # list of (N, 4)
        self.voxel_size = voxel_size
        assert len(self.descriptions) == len(self.voxels_sparse), \
            "描述与体素数量不一致，检查数据。"

    def __len__(self):
        return len(self.descriptions)

    def voxel_list_to_dense(self, voxel_ids):
        volume = torch.full((self.voxel_size, self.voxel_size, self.voxel_size),
                            EMPTY_BLOCK_ID, dtype=torch.long)
        for block_id, x, y, z in voxel_ids:
            if 0 <= x < self.voxel_size and 0 <= y < self.voxel_size and 0 <= z < self.voxel_size:
                volume[x, y, z] = block_id
        return volume

    def __getitem__(self, idx):
        desc = self.descriptions[idx]         # 文本描述 (string)
        voxel_ids = self.voxels_sparse[idx]   # 稀疏体素坐标列表

        # 稀疏 -> 稠密
        voxel_dense = self.voxel_list_to_dense(voxel_ids)
        
        return {
            "text": desc,
            "voxels": voxel_dense  # shape [64, 64, 64]
        }


################################
# 主训练脚本
################################
def train_model():
    import os
    import torch

    from torch.utils.data import DataLoader
    from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
    from shap_e.models.download import load_config

    # ========== 超参数 & 路径 ==========
    output_folder = 'processed_output'
    descriptions_path = os.path.join(output_folder, 'descriptions.pt')
    voxels_path = os.path.join(output_folder, 'voxels.pt')
    batch_size = 2
    num_epochs = 10

    # ========== 1. 加载数据集 ==========
    dataset = MinecraftDataset(descriptions_path, voxels_path, voxel_size=VOXEL_SIZE)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device:", device)

    # ========== 2. 初始化模型 ==========
    text_encoder = MinecraftCLIPTextEncoder(out_dim=512).to(device)
    voxel_encoder = VoxelEncoder(d_latent=512).to(device)
    voxel_decoder = VoxelDecoder(d_latent=512).to(device)

    # ========== 3. 加载并配置扩散模型 ==========
    # 这里以 shap-e 内部的 "text300M" 模型配置为例，你需要根据需求修改
    config = load_config("text300M", device=device)
    # 根据我们自定义的大小做适配
    config["inner"]["input_channels"] = 512
    config["inner"]["n_ctx"] = 1
    config["inner"]["output_channels"] = 512 * 2

    diffusion_model = diffusion_from_config(config).to(device)

    # ========== 4. 优化器配置 ==========
    # 如果想让 text_encoder 一起训练，可以把它的参数也加进来
    params = list(diffusion_model.parameters())
    params += list(voxel_encoder.parameters())
    params += list(voxel_decoder.parameters())
    # 也可以加：params += list(text_encoder.parameters())

    optimizer = optim.AdamW(params, lr=1e-4)

    # ========== 5. 训练循环 ==========
    for epoch in range(num_epochs):
        for step, batch in enumerate(dataloader):
            optimizer.zero_grad()

            # 取文本和体素
            text_list = batch["text"]  # list of strings
            voxels_dense = batch["voxels"]  # [batch_size, 64, 64, 64]

            # 转到 GPU
            voxels_dense = voxels_dense.to(device)
            
            # (batch, 64, 64, 64) -> (batch, 1, 64, 64, 64)
            voxels_input = voxels_dense.unsqueeze(1).float()

            batch_size_now = voxels_input.shape[0]

            # 生成时间步 t
            t = torch.randint(0, diffusion_model.num_timesteps, (batch_size_now,), device=device)

            # 编码体素
            encoded_voxels = voxel_encoder(voxels_input)  # [batch_size, 512]

            # 编码文本
            text_embeddings = text_encoder(text_list)     # [batch_size, 512]
            # 在这里如果扩散模型需要 (batch_size, 1, hidden_dim) 格式，可扩展一维
            text_embeddings = text_embeddings.unsqueeze(1)  # [batch_size, 1, 512]

            # 计算扩散模型训练损失
            terms = diffusion_model.training_losses(
                model=diffusion_model.wrapped,
                x_start=encoded_voxels,  # 初始潜变量
                t=t,
                model_kwargs={"cond": text_embeddings}
            )
            loss = terms["loss"].mean()

            loss.backward()
            optimizer.step()

            if step % 10 == 0:
                print(f"Epoch {epoch}, Step {step}, Loss: {loss.item():.4f}")

        # ========== 每个 epoch 结束后，可选地保存模型 & 采样 ==========
        torch.save(diffusion_model.state_dict(), os.path.join(output_folder, f'diffusion_model_epoch_{epoch}.pt'))
        torch.save(voxel_encoder.state_dict(),   os.path.join(output_folder, f'voxel_encoder_epoch_{epoch}.pt'))
        torch.save(voxel_decoder.state_dict(),   os.path.join(output_folder, f'voxel_decoder_epoch_{epoch}.pt'))
        torch.save(text_encoder.state_dict(),    os.path.join(output_folder, f'text_encoder_epoch_{epoch}.pt'))

        # 在此进行一次采样测试
        sample_texts = ["a redstone torch"]  # 你想要测试的文本
        with torch.no_grad():
            sample_text_embeddings = text_encoder(sample_texts).to(device)  # [1, 512]
            sample_text_embeddings = sample_text_embeddings.unsqueeze(1)    # [1, 1, 512]

            # 扩散采样 latent 向量
            sampled_latents = diffusion_model.sample_latents(
                batch_size=1,
                model=diffusion_model.wrapped,
                diffusion=diffusion_model,
                model_kwargs={"cond": sample_text_embeddings},
                guidance_scale=1.0,   # 超参数，可尝试调整
                clip_denoised=False,
                use_fp16=False,
                use_karras=True,
                karras_steps=64,
                sigma_min=diffusion_model.sigma_min,
                sigma_max=diffusion_model.sigma_max,
                s_churn=0,
                device=device,
                progress=True
            )
            # 解码体素
            sampled_voxels_logits = voxel_decoder(sampled_latents)     # [1, 600, 64, 64, 64]
            sampled_voxels = torch.argmax(sampled_voxels_logits, dim=1)  # [1, 64, 64, 64]
            sampled_voxels = sampled_voxels.squeeze(0).cpu().numpy()     # [64, 64, 64]

            # 将生成的体素保存到文本文件
            sample_voxel_text_lines = []
            for x in range(VOXEL_SIZE):
                for y in range(VOXEL_SIZE):
                    for z in range(VOXEL_SIZE):
                        block_id = sampled_voxels[x, y, z]
                        if block_id != EMPTY_BLOCK_ID:
                            sample_voxel_text_lines.append(f"{block_id} {x} {y} {z}")
            sample_voxel_text = '\n'.join(sample_voxel_text_lines)

            sample_voxel_txt_path = os.path.join(output_folder, f'sample_epoch_{epoch}.txt')
            with open(sample_voxel_txt_path, 'w', encoding='utf-8') as txt_file:
                txt_file.write(sample_voxel_text)

            print(f"[采样] Epoch {epoch} - 体素数据已保存到: {sample_voxel_txt_path}")

    print("训练完成。")


if __name__ == "__main__":
    train_model()

5. 采样与导出

上面的示例脚本在每个 epoch 结束后，都会对一段示例文本（如 "a redstone torch"）进行采样，然后：

1. 得到潜变量 sampled_latents
2. 用 voxel_decoder 得到体素分布 logits
3. 取 argmax 作为最终的方块 ID
4. 将非 EMPTY_BLOCK_ID 的坐标写回 txt 文件保存。

如果需要可视化或导入 Minecraft，你可以在生成的 .txt 文件基础上做进一步转换，比如：

• 直接写成 Minecraft .schematic 或者其他编辑器可识别的格式
• 在 Minecraft 中以指令方式构造这些坐标块

6. 完整示例脚本

为了方便使用，这里将数据预处理和训练都整合在一个脚本中（你也可以拆分为两个文件）。请根据你的实际需求修改路径、超参数、方块 ID 映射等内容。

python
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from transformers import CLIPTokenizer, CLIPTextModel

from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
from shap_e.models.download import load_config
from shap_e.models.generation.perceiver import SimplePerceiver
from shap_e.models.generation.transformer import Transformer

################################
# 0. 超参数和全局变量
################################
VOXEL_SIZE = 64        # 体素网格尺寸
EMPTY_BLOCK_ID = 0     # 空方块 ID
NUM_BLOCK_TYPES = 600  # 你需要根据实际 block_ids 数量修改

training_data_path = 'training_data/10001-1.jsonl'  # JSONL 数据
block_ids_path = 'blockids.txt'                     # block_ids 映射脚本
output_folder = 'processed_output'
os.makedirs(output_folder, exist_ok=True)

################################
# 1. 预处理数据 (JSONL -> .pt)
################################
def prepare_data():
    # 读取 block_ids
    block_ids = {}
    with open(block_ids_path, 'r', encoding='utf-8') as f:
        exec(f.read())  # 例如: block_ids = {0: "air", 1: "stone", ...}

    # 这里仅示例如何保存原始文本，不做提前tokenize
    description_list = []
    voxel_list = []

    with open(training_data_path, 'r', encoding='utf-8') as f:
        for line in f:
            data = json.loads(line.strip())
            desc = data.get('input', '')
            voxel_data = data.get('output', '')

            # 1) 保存描述
            description_list.append(desc)

            # 2) 解析体素
            lines = voxel_data.strip().split('\n')
            sparse_voxels = []
            for l in lines:
                parts = l.strip().split()
                if len(parts) == 4:
                    block_name, x, y, z = parts
                    # 根据名称找 ID
                    b_id = None
                    for k, v in block_ids.items():
                        if v == block_name:
                            b_id = k
                            break
                    if b_id is not None:
                        sparse_voxels.append([int(b_id), int(x), int(y), int(z)])
                    else:
                        print(f"[警告] 未找到 block_name='{block_name}' 的映射，已跳过。")
                else:
                    print(f"[警告] 无效体素行: {l}")
            voxel_list.append(sparse_voxels)

    torch.save(description_list, os.path.join(output_folder, 'descriptions.pt'))
    torch.save(voxel_list, os.path.join(output_folder, 'voxels.pt'))
    print("数据预处理完成。")


################################
# 2. 定义数据集和模型
################################
class MinecraftDataset(Dataset):
    def __init__(self, descriptions_path, voxels_path, voxel_size=64):
        super().__init__()
        self.descriptions = torch.load(descriptions_path)  # list of str
        self.voxels_sparse = torch.load(voxels_path)       # list of (N,4)
        self.voxel_size = voxel_size
        assert len(self.descriptions) == len(self.voxels_sparse), \
            "描述与体素数量不一致。"

    def __len__(self):
        return len(self.descriptions)

    def voxel_list_to_dense(self, voxel_ids):
        volume = torch.full((self.voxel_size, self.voxel_size, self.voxel_size),
                            EMPTY_BLOCK_ID, dtype=torch.long)
        for (block_id, x, y, z) in voxel_ids:
            if 0 <= x < self.voxel_size and 0 <= y < self.voxel_size and 0 <= z < self.voxel_size:
                volume[x, y, z] = block_id
        return volume

    def __getitem__(self, idx):
        desc = self.descriptions[idx]
        sparse_voxels = self.voxels_sparse[idx]
        dense_voxel = self.voxel_list_to_dense(sparse_voxels)
        return {
            "text": desc,
            "voxels": dense_voxel
        }

class MinecraftCLIPTextEncoder(nn.Module):
    def __init__(self, clip_model_name="openai/clip-vit-large-patch14", out_dim=512):
        super().__init__()
        self.clip = CLIPTextModel.from_pretrained(clip_model_name)
        self.tokenizer = CLIPTokenizer.from_pretrained(clip_model_name)
        self.proj = nn.Linear(self.clip.config.hidden_size, out_dim)

    def forward(self, text_list):
        inputs = self.tokenizer(text_list, return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(self.clip.device) for k, v in inputs.items()}
        outputs = self.clip(**inputs)
        pooled = outputs.pooler_output
        return self.proj(pooled)

class VoxelEncoder(nn.Module):
    def __init__(self, d_latent=512):
        super().__init__()
        self.conv1 = nn.Conv3d(1, 32, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.Conv3d(32, 64, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.Conv3d(64, 128, kernel_size=4, stride=2, padding=1)
        # 注意卷积堆叠后空间会缩小到 8x8x8(64/2/2/2 = 8)
        self.fc = nn.Linear(128 * 8 * 8 * 8, d_latent)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

class VoxelDecoder(nn.Module):
    def __init__(self, d_latent=512, output_size=64, n_channels=NUM_BLOCK_TYPES):
        super().__init__()
        self.fc = nn.Linear(d_latent, 128 * 8 * 8 * 8)
        self.deconv1 = nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, padding=1)
        self.deconv2 = nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, padding=1)
        self.deconv3 = nn.ConvTranspose3d(32, n_channels, kernel_size=4, stride=2, padding=1)

    def forward(self, x):
        x = self.fc(x)
        x = x.view(x.shape[0], 128, 8, 8, 8)
        x = F.relu(self.deconv1(x))
        x = F.relu(self.deconv2(x))
        x = self.deconv3(x)  # (batch, n_channels, 64, 64, 64)
        return x

################################
# 3. 训练主函数
################################
def train_model():
    # 1) 先构造 DataLoader
    dataset = MinecraftDataset(
        os.path.join(output_folder, 'descriptions.pt'),
        os.path.join(output_folder, 'voxels.pt'),
        voxel_size=VOXEL_SIZE
    )
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

    # 2) 初始化设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("[Info] Using device:", device)

    # 3) 初始化模型
    text_encoder = MinecraftCLIPTextEncoder(out_dim=512).to(device)
    voxel_encoder = VoxelEncoder(d_latent=512).to(device)
    voxel_decoder = VoxelDecoder(d_latent=512).to(device)

    # 4) 加载扩散模型配置
    config = load_config("text300M", device=device)
    config["inner"]["input_channels"] = 512
    config["inner"]["n_ctx"] = 1
    config["inner"]["output_channels"] = 512 * 2
    diffusion_model = diffusion_from_config(config).to(device)

    # 5) 优化器
    params = list(diffusion_model.parameters()) \
           + list(voxel_encoder.parameters()) \
           + list(voxel_decoder.parameters())
    # 如果需要训练文本编码器则一起加上: params += list(text_encoder.parameters())
    optimizer = optim.AdamW(params, lr=1e-4)

    num_epochs = 5
    for epoch in range(num_epochs):
        for step, batch in enumerate(dataloader):
            optimizer.zero_grad()

            text_list = batch["text"]
            voxels_dense = batch["voxels"].to(device)  # (B, 64, 64, 64)
            voxels_dense = voxels_dense.unsqueeze(1).float()  # (B,1,64,64,64)

            B = voxels_dense.size(0)
            t = torch.randint(0, diffusion_model.num_timesteps, (B,), device=device)

            # 编码
            encoded_voxels = voxel_encoder(voxels_dense)   # (B, 512)
            text_emb = text_encoder(text_list)             # (B, 512)
            text_emb = text_emb.unsqueeze(1)               # (B,1,512)

            # 扩散损失
            terms = diffusion_model.training_losses(
                model=diffusion_model.wrapped,
                x_start=encoded_voxels,
                t=t,
                model_kwargs={"cond": text_emb}
            )
            loss = terms["loss"].mean()
            loss.backward()
            optimizer.step()

            if step % 10 == 0:
                print(f"[Epoch {epoch}/{num_epochs}] Step {step}, Loss={loss.item():.4f}")

        # ========== 每个 epoch 做一次保存和采样 ==========
        torch.save(diffusion_model.state_dict(), os.path.join(output_folder, f'diffusion_model_e{epoch}.pt'))
        torch.save(voxel_encoder.state_dict(),   os.path.join(output_folder, f'voxel_encoder_e{epoch}.pt'))
        torch.save(voxel_decoder.state_dict(),   os.path.join(output_folder, f'voxel_decoder_e{epoch}.pt'))
        torch.save(text_encoder.state_dict(),    os.path.join(output_folder, f'text_encoder_e{epoch}.pt'))

        # 测试采样
        with torch.no_grad():
            sample_text = ["a redstone torch"]
            sample_emb = text_encoder(sample_text).to(device)  # (1,512)
            sample_emb = sample_emb.unsqueeze(1)               # (1,1,512)

            # 采样 latent
            sampled_latents = diffusion_model.sample_latents(
                batch_size=1,
                model=diffusion_model.wrapped,
                diffusion=diffusion_model,
                model_kwargs={"cond": sample_emb},
                guidance_scale=1.0,
                clip_denoised=False,
                use_fp16=False,
                use_karras=True,
                karras_steps=64,
                sigma_min=diffusion_model.sigma_min,
                sigma_max=diffusion_model.sigma_max,
                s_churn=0,
                device=device,
                progress=True
            )
            # 解码体素
            logits = voxel_decoder(sampled_latents)         # (1, NUM_BLOCK_TYPES,64,64,64)
            pred_blocks = torch.argmax(logits, dim=1)[0]    # (64,64,64)

            # 输出到txt
            lines = []
            arr = pred_blocks.cpu().numpy()
            for x in range(VOXEL_SIZE):
                for y in range(VOXEL_SIZE):
                    for z in range(VOXEL_SIZE):
                        bid = arr[x, y, z]
                        if bid != EMPTY_BLOCK_ID:
                            lines.append(f"{bid} {x} {y} {z}")

            save_path = os.path.join(output_folder, f"sample_e{epoch}.txt")
            with open(save_path, "w", encoding="utf-8") as fw:
                fw.write("\n".join(lines))
            print(f"[采样] 生成体素已保存至 {save_path}")

    print("[Info] 训练全部完成。")


if __name__ == "__main__":
    # 第一步：预处理数据
    prepare_data()
    # 第二步：训练模型
    train_model()

运行本脚本后，将依次完成：

1. 数据预处理：从 JSONL 读取并解析为 descriptions.pt & voxels.pt
2. 模型训练：读入数据集、初始化网络与扩散模型，在多个 epoch 中进行优化
3. 周期性采样：每个 epoch 完成后，使用示例文本进行一次 3D 体素采样并导出到 .txt 文件

注意：脚本中关于 方块种类（NUM_BLOCK_TYPES）、体素分辨率、网络结构、超参数（学习率、batch size 等） 都是示例值，需要你根据实际数据规模和硬件资源进行修改与试验。

结论

通过以上示例脚本，你就可以微调（或训练）一个 Shap-E 风格的“文本->隐式->体素”生成流程。其中有几个关键要点需要在实际项目中关注：

1. 数据质量与多样性：Minecraft 体素往往需要大量且多样化的数据来训练出具有泛化能力的模型。
2. 合适的网络结构：如果你需要更高分辨率或更复杂的形状，可以加深卷积网络或采用其他 3D 编码结构（如 3D Swin Transformer 等）。
3. 扩散模型调参：包括时间步数、guidance_scale、use_karras 等超参数，都可能影响最终生成质量和效率。
4. 推理输出后处理：将生成的体素文件导入 Minecraft 时，需要再做坐标转换、方块边界检测或使用相应的世界编辑器格式（WorldEdit / Schematic 等）。

以上示例可作为一个“端到端”的起点，后续你可以根据具体需求自定义或改进各个模块。祝你在将 Shap-E 应用于 Minecraft 体素生成的过程中一切顺利!