Attempting to solve pickle issue

demo.py

@@ -23,8 +23,10 @@ from mast3r.utils.misc import hash_md5
 import main
 import utils.export as export
 
-@spaces.GPU(duration=10)
-def get_reconstructed_scene(outdir, model, device, silent, image_size, ios_mode, filelist):
+@spaces.GPU(duration=15)
+def get_reconstructed_scene(outdir, model, silent, image_size, ios_mode, filelist):
+
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
  assert len(filelist) == 1 or len(filelist) == 2, "Please provide one or two images"
  if ios_mode:
@@ -37,6 +39,7 @@ def get_reconstructed_scene(outdir, model, device, silent, image_size, ios_mode,
  img['img'] = img['img'].to(device)
  img['original_img'] = img['original_img'].to(device)
  img['true_shape'] = torch.from_numpy(img['true_shape'])
+ model = model.to(device)
 
  output = model(imgs[0], imgs[1])
 
@@ -50,13 +53,11 @@ if __name__ == '__main__':
  image_size = 512
  silent = False
  ios_mode = True
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
  model_name = "brandonsmart/splatt3r_v1.0"
  filename = "epoch=19-step=1200.ckpt"
  weights_path = hf_hub_download(repo_id=model_name, filename=filename)
  model = main.MAST3RGaussians.load_from_checkpoint(weights_path, 'cpu')
- model = model.to(device)
  chkpt_tag = hash_md5(weights_path)
 
  # Define example inputs and their corresponding precalculated outputs
@@ -88,7 +89,7 @@ if __name__ == '__main__':
  cache_path = os.path.join(tmpdirname, chkpt_tag)
  os.makedirs(cache_path, exist_ok=True)
 
- recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, model, device, silent, image_size, ios_mode)
+ recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, model, silent, image_size, ios_mode)
 
  if not ios_mode:
  for i in range(len(examples)):