Update app.py

app.py

@@ -256,7 +256,7 @@ def add_file(history, task_history, file, task_new, fig, query):
         coords        = ((torch.from_numpy(position_recon) + 0.5) * 64).int().contiguous()
         ss            = torch.zeros(1, 64, 64, 64, dtype=torch.long)
         ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
-        token         = vqvae.Encode(ss.to(dtype=torch.float32).unsqueeze(0).to("cuda"))
+        token         = vqvae.Encode(ss.to(dtype=torch.float32).unsqueeze(0))
         token         = token[0].cpu().numpy().tolist()
         words         = token_to_words(token)
         fig            = make_pointcloud_figure(position_recon,rotate=True)
@@ -387,9 +387,8 @@ def _transform_messages(original_messages):
     return transformed_messages
 
 from trellis.models.sparse_structure_vqvae import VQVAE3D
-device       = torch.device("cuda")
+device       = torch.device("cpu")
 vqvae        = VQVAE3D(num_embeddings=8192)
-device       = torch.device("cuda")
 vqvae.eval()
 filepath = hf_hub_download(repo_id="yejunliang23/3DVQVAE",filename="3DVQVAE.bin")
 state_dict = torch.load(filepath, map_location="cpu")
@@ -398,7 +397,7 @@ vqvae=vqvae.to(device)
 
 MODEL_DIR = "yejunliang23/ShapeLLM-7B-omni"
 model_ckpt_path=MODEL_DIR
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_ckpt_path, torch_dtype="auto", device_map={"": 0})
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_ckpt_path, torch_dtype="auto", device_map="auto")
 processor = AutoProcessor.from_pretrained(model_ckpt_path)
 tokenizer = processor.tokenizer
 from huggingface_hub import hf_hub_download