yejunliang23 commited on
Commit
deccfda
·
verified ·
1 Parent(s): 5db9697

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -64,7 +64,8 @@ def save_ply_from_array(verts):
64
  tmpf.flush(); tmpf.close()
65
  return tmpf.name
66
 
67
- @spaces.GPU(duration=120)
 
68
  def predict(_chatbot,task_history,viewer_voxel,viewer_mesh,task_new,seed,top_k,top_p,temperature):
69
  torch.manual_seed(seed)
70
  chat_query = _chatbot[-1][0]
@@ -393,6 +394,10 @@ def _transform_messages(original_messages):
393
 
394
  return transformed_messages
395
 
 
 
 
 
396
  from trellis.models.sparse_structure_vqvae import VQVAE3D
397
  device = torch.device("cuda")
398
  vqvae = VQVAE3D(num_embeddings=8192)
 
64
  tmpf.flush(); tmpf.close()
65
  return tmpf.name
66
 
67
+ #@spaces.GPU(duration=120)
68
+ @spaces.GPU
69
  def predict(_chatbot,task_history,viewer_voxel,viewer_mesh,task_new,seed,top_k,top_p,temperature):
70
  torch.manual_seed(seed)
71
  chat_query = _chatbot[-1][0]
 
394
 
395
  return transformed_messages
396
 
397
+ print(f"CUDA Available: {torch.cuda.is_available()}")
398
+ print(f"CUDA Version: {torch.version.cuda}")
399
+ print(f"Number of GPUs: {torch.cuda.device_count()}")
400
+
401
  from trellis.models.sparse_structure_vqvae import VQVAE3D
402
  device = torch.device("cuda")
403
  vqvae = VQVAE3D(num_embeddings=8192)