Trellis 2 se ejecuta con éxito en ROCm 7.11 con AMD RX 9070 XT.

Haciendo funcionar Trellis 2 en hardware AMD

Un desarrollador ha ejecutado exitosamente Trellis 2 en una GPU AMD RX 9070 XT usando ROCm 7.11 en Linux Mint 22.3. Esto aborda problemas comunes donde los usuarios encontraban corte de geometría, fallos en la vista previa y otros errores al intentar ejecutar Trellis 2 en hardware AMD.

Problemas clave y soluciones

El desarrollador identificó dos problemas principales que causaban la mayoría de los fallos:

1. Inestabilidad de ROCm con tensores N altos

Las operaciones de ROCm se vuelven inestables con tensores grandes, causando desbordamientos o valores NaN. El código original en linear.py en la carpeta sparse utilizaba:

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

La solución implementa procesamiento por fragmentos para evitar problemas con ROCm:

ROCM_SAFE_CHUNK = 524_288
def rocm_safe_linear(feats: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor:
    """F.linear con solución alternativa de fragmentación para N grande en ROCm."""
    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. Función hipMemcpy2D defectuosa en CuMesh

La función hipMemcpy2D en CuMesh causaba que los vértices y caras desaparecieran o se corrompieran. La inicialización original de CuMesh utilizaba:

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
    ));
    ...
}

La solución reemplaza la copia 2D con una versión 1D:

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

Resultados y rendimiento

Con estas soluciones, el desarrollador logró que el flujo de trabajo de imagen-a-3D funcionara, incluyendo la renderización de vista previa (sin normales) y la exportación final GLB. En una imagen de prueba con 21,204 tokens, el proceso tomó aproximadamente 280 segundos desde el inicio hasta la generación de la vista previa. La ejecución utilizó resolución 1024 con todos los muestreadores configurados a 20 pasos.