global model

app.py

@@ -21,7 +21,7 @@ print(ckpt_path)
 # else:
 #     mark_dtype = torch.float16
 #     device = torch.device('cpu')
-mark_dtype = torch.float16
+mark_dtype = torch.bfloat16
 device = torch.device('cuda')
 
 model = VGGT_OriAny_Ref(out_dim=900, dtype=mark_dtype, nopretrain=True)
@@ -69,7 +69,8 @@ def run_inference(image_ref, image_tgt, do_rm_bkg):
             pil_ref = background_preprocess(pil_ref, True)
 
     try:
-        ans_dict = inf_single_case(model, pil_ref, pil_tgt)
+        # ans_dict = inf_single_case(model, pil_ref, pil_tgt)
+        ans_dict = inf_single_case(pil_ref, pil_tgt)
     except Exception as e:
         print("Inference error:", e)
         raise gr.Error(f"Inference failed: {str(e)}")

inference.py

@@ -177,7 +177,8 @@ def preprocess_images(image_list, mode="crop"):
     return images
 
 @torch.no_grad()
-def inf_single_batch(model, batch):
+def inf_single_batch(batch):
+    global model
     device = model.get_device()
     batch_img_inputs = batch # (B, S, 3, H, W)
     # print(batch_img_inputs.shape)
@@ -229,12 +230,13 @@ def inf_single_batch(model, batch):
 # input PIL Image
 @spaces.GPU
 @torch.no_grad()
-def inf_single_case(model, image_ref, image_tgt):
+def inf_single_case(image_ref, image_tgt):
+    global model
     if image_tgt is None:
         image_list = [image_ref]
     else:
         image_list = [image_ref, image_tgt]
     image_tensors = preprocess_images(image_list, mode="pad").to('cuda')
-    ans_dict = inf_single_batch(model=model, batch=image_tensors.unsqueeze(0))
+    ans_dict = inf_single_batch(batch=image_tensors.unsqueeze(0))
     print(ans_dict)
     return ans_dict