search

iffyloop / TripoSR-Bake

Generate 3D meshes from a single 2D image using TripoSR, complete with manual geometry editing and texture baking support
The Unlicense
50 stars 0 forks source link

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! #7

Open videoprice opened 4 months ago

videoprice commented 4 months ago

Hello, have this error both cpu and cuda

(1/5): Load TripoSR and scene codes (2/5): Load mesh (3/5): Generate UVs (4/5): Rasterize UV atlas (5/5): Sample NeRF to UV atlas Traceback (most recent call last): File "D:\TripoSR-Bake\02-texture.py", line 268, in main() File "D:\TripoSR-Bake\02-texture.py", line 247, in main bake_result = run_bake(args, tsr_result, rasterize_result) File "D:\TripoSR-Bake\02-texture.py", line 167, in run_bake queried_grid = tsr_result["model"].renderer.query_triplane( File "D:\TripoSR-Bake\tsr\models\nerf_renderer.py", line 78, in query_triplane net_out = chunk_batch(_query_chunk, self.chunk_size, positions) File "D:\TripoSR-Bake\tsr\utils.py", line 167, in chunk_batch out_chunk = func( File "D:\TripoSR-Bake\tsr\models\nerf_renderer.py", line 61, in _query_chunk out: torch.Tensor = F.grid_sample( File "D:\TripoSR-Bake\venv\lib\site-packages\torch\nn\functional.py", line 4351, 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)

leeooo002 commented 4 months ago

the same problem ^---^

jorgeolothar commented 3 months ago

Seems like when scene_codes is loaded from the .pkl file, it can no longer be moved between devices. You can add this code underneath 02-texture.py -> scene_codes = pickle.load(infile) to recreate it.

        # Convert to numpy and back to tensor. Unpickled tensor will not change devices for some reason.
        np_arr = scene_codes.cpu().numpy()
        scene_codes = torch.from_numpy(np_arr)
        # Move model and tensor to CPU
        model.cpu()
        scene_codes.cpu()

The python file will run with a --device different from cpu. But it will do the texturing there. If you want to run it in the gpu, the model, scene_codes, and a bunch of other intermediate tensors in nerf_renderer.py need to be moved to the gpu as well.

iffyloop commented 1 month ago

Hey everyone, sorry I haven't been active on this repository for a couple of months. Thank you @jorgeolothar for posting your solution! I'm planning to integrate this into a separate project eventually and then archive this repo, but hopefully in the meantime people will see this thread and apply the fix themselves.