--bake-texture option doesn't work with gpu, here is the error message:

[Andy Triboletti]

--bake-texture option doesn't work with gpu. without that option it works. here is the error message:

(newtest) andy@andys-pc:/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR$ python run.py featured_image.png --output-dir output/
--bake-texture
2024-05-23 10:59:57,568 - INFO - Initializing model ...
2024-05-23 11:00:01,162 - INFO - Initializing model finished in 3594.59ms.
2024-05-23 11:00:01,162 - INFO - Processing images ...
2024-05-23 11:00:01,875 - INFO - Processing images finished in 712.37ms.
2024-05-23 11:00:01,875 - INFO - Running image 1/1 ...
2024-05-23 11:00:01,875 - INFO - Running model ...
2024-05-23 11:00:02,659 - INFO - Running model finished in 784.47ms.
2024-05-23 11:00:02,660 - INFO - Extracting mesh ...
2024-05-23 11:00:04,841 - INFO - Extracting mesh finished in 2181.52ms.
2024-05-23 11:00:04,841 - INFO - Baking texture ...
WARNING: dzn is not a conformant Vulkan implementation, testing use only.
WARNING: dzn is not a conformant Vulkan implementation, testing use only.
WARNING: Some incorrect rendering might occur because the selected Vulkan device (Microsoft Direct3D12 (NVIDIA GeForce RTX 3060)) doesn't support base Zink requirements: feats.features.logicOp have_EXT_custom_border_color have_EXT_line_rasterization
glx: failed to create drisw screen
failed to load driver: zink
Traceback (most recent call last):
  File "/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR/run.py", line 181, in <module>
    bake_output = bake_texture(meshes[0], model, scene_codes[0], args.texture_resolution)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR/tsr/bake_texture.py", line 162, in bake_texture
    colors_texture = positions_to_colors(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR/tsr/bake_texture.py", line 140, in positions_to_colors
    queried_grid = model.renderer.query_triplane(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR/tsr/models/nerf_renderer.py", line 78, in query_triplane
    net_out = chunk_batch(_query_chunk, self.chunk_size, positions)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR/tsr/utils.py", line 169, in chunk_batch
    out_chunk = func(
                ^^^^^
  File "/mnt/c/Users/andy/OneDrive/Documents/GitHub/TripoSR/tsr/models/nerf_renderer.py", line 61, in _query_chunk
    out: torch.Tensor = F.grid_sample(
                        ^^^^^^^^^^^^^^
  File "/home/andy/miniconda3/envs/newtest/lib/python3.11/site-packages/torch/nn/functional.py", line 4244, in grid_sample
    return torch.grid_sampler(input, grid, mode_enum, padding_mode_enum, align_corners)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument grid in method wrapper_CUDA__grid_sampler_2d)

[Andy Triboletti]

when I add --device cpu it generates the texture file.

[Andy Triboletti]

when I try --device cuda it fails

[Andy Triboletti]

Using this command works:
python run.py featured_image.png --output-dir output/ --device cuda:0 --model-save-format glb

I used --device cuda:0 instead of --device cuda.

But not when using .obj, only cpu works there.

[Luca Bonfiglioli]

Replace positions_to_colors in bake_texture.py line 137 with the following code:

def positions_to_colors(model, scene_code, positions_texture, texture_resolution):
    positions = torch.tensor(
        positions_texture.reshape(-1, 4)[:, :-1], device=scene_code.device
    )
    with torch.no_grad():
        queried_grid = model.renderer.query_triplane(
            model.decoder,
            positions,
            scene_code,
        )
    rgb_f = queried_grid["color"].cpu().numpy().reshape(-1, 3)
    rgba_f = np.insert(rgb_f, 3, positions_texture.reshape(-1, 4)[:, -1], axis=1)
    rgba_f[rgba_f[:, -1] == 0.0] = [0, 0, 0, 0]
    return rgba_f.reshape(texture_resolution, texture_resolution, 4)