Trellis 2 успешно работает на ROCm 7.11 с видеокартой AMD RX 9070 XT.

Запуск Trellis 2 на оборудовании AMD

Разработчик успешно запустил Trellis 2 на видеокарте AMD RX 9070 XT с использованием ROCm 7.11 на Linux Mint 22.3. Это решает распространённые проблемы, с которыми сталкивались пользователи при попытке запуска Trellis 2 на оборудовании AMD, такие как обрезка геометрии, сбои предпросмотра и другие ошибки.

Ключевые проблемы и решения

Разработчик определил две основные проблемы, которые вызывали большинство сбоев:

1. Нестабильность ROCm с тензорами высокого порядка N

Операции ROCm становятся нестабильными с большими тензорами, вызывая переполнения или значения NaN. Исходный код в файле linear.py в папке sparse использовал:

def forward(self, input: VarLenTensor) -> VarLenTensor:
    return input.replace(super().forward(input.feats))

Исправление реализует обработку по частям, чтобы избежать проблем с ROCm:

ROCM_SAFE_CHUNK = 524_288
def rocm_safe_linear(feats: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor:
    """F.linear с обходным решением для ROCm с разбиением на части при большом N."""
    N = feats.shape[0]
    if N <= ROCM_SAFE_CHUNK:
        return F.linear(feats, weight, bias)
    out = torch.empty(N, weight.shape[0], device=feats.device, dtype=feats.dtype)
    for s in range(0, N, ROCM_SAFE_CHUNK):
        e = min(s + ROCM_SAFE_CHUNK, N)
        out[s:e] = F.linear(feats[s:e], weight, bias)
    return out

def forward(self, input):
    feats = input.feats if hasattr(input, 'feats') else input
    out = rocm_safe_linear(feats, self.weight, self.bias)
    if hasattr(input, 'replace'):
        return input.replace(out)
    return out

2. Неисправная функция hipMemcpy2D в CuMesh

Функция hipMemcpy2D в CuMesh вызывала пропадание или повреждение вершин и граней. Исходная инициализация CuMesh использовала:

void CuMesh::init(const torch::Tensor& vertices, const torch::Tensor& faces) {
    size_t num_vertices = vertices.size(0);
    size_t num_faces = faces.size(0);
    this->vertices.resize(num_vertices);
    this->faces.resize(num_faces);
    CUDA_CHECK(cudaMemcpy2D(
        this->vertices.ptr,
        sizeof(float3),
        vertices.data_ptr(),
        sizeof(float) * 3,
        sizeof(float) * 3,
        num_vertices,
        cudaMemcpyDeviceToDevice
    ));
    ...
}

Исправление заменяет 2D-копирование на 1D-версию:

CUDA_CHECK(cudaMemcpy(
    this->vertices.ptr,
    vertices.data_ptr(),
    num_vertices * sizeof(float3),
    cudaMemcpyDeviceToDevice
));

Результаты и производительность

С этими исправлениями разработчик успешно заставил работать конвейер преобразования изображения в 3D, включая предварительный рендеринг (без нормалей) и окончательный экспорт в GLB. На тестовом изображении с 21 204 токенами процесс занял примерно 280 секунд от начала до генерации предпросмотра. Запуск использовал разрешение 1024 со всеми сэмплерами, установленными на 20 шагов.