cuda

app.py

@@ -15,11 +15,12 @@ from huggingface_hub import hf_hub_download
 ckpt_path = hf_hub_download(repo_id=ORIANY_V2, filename=REMOTE_CKPT_PATH, repo_type="model", cache_dir='./', resume_download=True)
 print(ckpt_path)
 
-
-# mark_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
-mark_dtype = torch.float16
-# device = 'cuda:0'
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.cuda.is_available():
+    mark_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+    device = torch.device('cuda')
+else:
+    mark_dtype = torch.float16
+    device = torch.device('cpu')
 
 model = VGGT_OriAny_Ref(out_dim=900, dtype=mark_dtype, nopretrain=True)
 model.load_state_dict(torch.load(ckpt_path, map_location='cpu'))
@@ -143,8 +144,8 @@ def run_inference(image_ref, image_tgt, do_rm_bkg):
 
 
 # ====== Gradio Blocks UI ======
-with gr.Blocks(title="Orient-Anything Demo") as demo:
-    gr.Markdown("# Orient-Anything Demo")
+with gr.Blocks(title="Orient-Anything-V2 Demo") as demo:
+    gr.Markdown("# Orient-Anything-V2 Demo")
     gr.Markdown("Upload a **reference image** (required). Optionally upload a **target image** for relative pose.")
 
     with gr.Row():