Update app.py

app.py

@@ -21,10 +21,6 @@ from huggingface_hub import hf_hub_download
 import numpy as np
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-import torch
-print("CUDA available:", torch.cuda.is_available())
-print("CUDA device count:", torch.cuda.device_count())
-1/0
 def _remove_image_special(text):
     text = text.replace('<ref>', '').replace('</ref>', '')
     return re.sub(r'<box>.*?(</box>|$)', '', text)
@@ -106,7 +102,7 @@ def predict(_chatbot,task_history,viewer_voxel,viewer_mesh,task_new,seed,top_k,t
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor(text=[text], images=image_inputs,videos=video_inputs, padding=True, return_tensors='pt')
-    inputs = inputs.to(model.device)
+    inputs = inputs.to("cuda")
 
     eos_token_id = [tokenizer.eos_token_id,159858]
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
@@ -134,7 +130,7 @@ def predict(_chatbot,task_history,viewer_voxel,viewer_mesh,task_new,seed,top_k,t
 
     if encoding_indices is not None:
         print("processing mesh...")
-        recon = vqvae.Decode(encoding_indices.to(model.device))
+        recon = vqvae.Decode(encoding_indices.to("cuda"))
         z_s           = recon[0].detach().cpu() 
         z_s           = (z_s>0)*1      
         indices       = torch.nonzero(z_s[0] == 1)  
@@ -148,7 +144,7 @@ def predict(_chatbot,task_history,viewer_voxel,viewer_mesh,task_new,seed,top_k,t
         ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
         ss=ss.unsqueeze(0)
         coords = torch.argwhere(ss>0)[:, [0, 2, 3, 4]].int()
-        coords = coords.to(model.device)
+        coords = coords.to("cuda")
         try:
             print("processing mesh...")
             if len(image_lst) == 0: