fix NameError

app.py

@@ -10,6 +10,7 @@ from vision_tower import VGGT_OriAny_Ref
 from inference import *
 from app_utils import *
 from axis_renderer import BlendRenderer
+import spaces
 
 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)
@@ -32,6 +33,69 @@ print('Model loaded.')
 
 axis_renderer = BlendRenderer(RENDER_FILE)
 
+@spaces.GPU
+@torch.no_grad()
+def inf_single_batch(batch):
+    device = model.get_device()
+    batch_img_inputs = batch # (B, S, 3, H, W)
+    # print(batch_img_inputs.shape)
+    B, S, C, H, W = batch_img_inputs.shape
+    pose_enc = model(batch_img_inputs) # (B, S, D) S = 1
+    
+    pose_enc = pose_enc.view(B*S, -1)
+    angle_az_pred = torch.argmax(pose_enc[:, 0:360]       , dim=-1)
+    angle_el_pred = torch.argmax(pose_enc[:, 360:360+180] , dim=-1) - 90
+    angle_ro_pred = torch.argmax(pose_enc[:, 360+180:360+180+360] , dim=-1) - 180
+    
+    # ori_val
+    # trained with BCE loss
+    distribute = F.sigmoid(pose_enc[:, 0:360]).cpu().float().numpy()
+    # trained with CE loss
+    # distribute = pose_enc[:, 0:360].cpu().float().numpy()
+    alpha_pred = val_fit_alpha(distribute = distribute)
+
+    # ref_val
+    if S > 1:
+        ref_az_pred = angle_az_pred.reshape(B,S)[:,0]
+        ref_el_pred = angle_el_pred.reshape(B,S)[:,0]
+        ref_ro_pred = angle_ro_pred.reshape(B,S)[:,0]
+        ref_alpha_pred = alpha_pred.reshape(B,S)[:,0]
+        rel_az_pred = angle_az_pred.reshape(B,S)[:,1]
+        rel_el_pred = angle_el_pred.reshape(B,S)[:,1]
+        rel_ro_pred = angle_ro_pred.reshape(B,S)[:,1]
+    else:
+        ref_az_pred = angle_az_pred[0]
+        ref_el_pred = angle_el_pred[0]
+        ref_ro_pred = angle_ro_pred[0]
+        ref_alpha_pred = alpha_pred[0]
+        rel_az_pred = 0.
+        rel_el_pred = 0.
+        rel_ro_pred = 0.
+
+    ans_dict = {
+        'ref_az_pred': ref_az_pred,
+        'ref_el_pred': ref_el_pred,
+        'ref_ro_pred': ref_ro_pred,
+        'ref_alpha_pred' : ref_alpha_pred,
+        'rel_az_pred'  : rel_az_pred,
+        'rel_el_pred'  : rel_el_pred,
+        'rel_ro_pred'  : rel_ro_pred,
+    }
+    
+    return ans_dict 
+
+# input PIL Image
+@torch.no_grad()
+def inf_single_case(image_ref, image_tgt):
+    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(batch=image_tensors.unsqueeze(0))
+    print(ans_dict)
+    return ans_dict
+
 
 # ====== 工具函数：安全图像处理 ======
 def safe_image_input(image):

inference.py

@@ -9,7 +9,7 @@ from scipy.special import i0
 from scipy.optimize import curve_fit
 from scipy.integrate import trapezoid
 from functools import partial
-import spaces
+
 
 def von_mises_pdf_alpha_numpy(alpha, x, mu, kappa):
     normalization = 2 * np.pi
@@ -175,68 +175,3 @@ def preprocess_images(image_list, mode="crop"):
             images = images.unsqueeze(0)
 
     return images
-
-@torch.no_grad()
-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)
-    B, S, C, H, W = batch_img_inputs.shape
-    pose_enc = model(batch_img_inputs) # (B, S, D) S = 1
-    
-    pose_enc = pose_enc.view(B*S, -1)
-    angle_az_pred = torch.argmax(pose_enc[:, 0:360]       , dim=-1)
-    angle_el_pred = torch.argmax(pose_enc[:, 360:360+180] , dim=-1) - 90
-    angle_ro_pred = torch.argmax(pose_enc[:, 360+180:360+180+360] , dim=-1) - 180
-    
-    # ori_val
-    # trained with BCE loss
-    distribute = F.sigmoid(pose_enc[:, 0:360]).cpu().float().numpy()
-    # trained with CE loss
-    # distribute = pose_enc[:, 0:360].cpu().float().numpy()
-    alpha_pred = val_fit_alpha(distribute = distribute)
-
-    # ref_val
-    if S > 1:
-        ref_az_pred = angle_az_pred.reshape(B,S)[:,0]
-        ref_el_pred = angle_el_pred.reshape(B,S)[:,0]
-        ref_ro_pred = angle_ro_pred.reshape(B,S)[:,0]
-        ref_alpha_pred = alpha_pred.reshape(B,S)[:,0]
-        rel_az_pred = angle_az_pred.reshape(B,S)[:,1]
-        rel_el_pred = angle_el_pred.reshape(B,S)[:,1]
-        rel_ro_pred = angle_ro_pred.reshape(B,S)[:,1]
-    else:
-        ref_az_pred = angle_az_pred[0]
-        ref_el_pred = angle_el_pred[0]
-        ref_ro_pred = angle_ro_pred[0]
-        ref_alpha_pred = alpha_pred[0]
-        rel_az_pred = 0.
-        rel_el_pred = 0.
-        rel_ro_pred = 0.
-
-    ans_dict = {
-        'ref_az_pred': ref_az_pred,
-        'ref_el_pred': ref_el_pred,
-        'ref_ro_pred': ref_ro_pred,
-        'ref_alpha_pred' : ref_alpha_pred,
-        'rel_az_pred'  : rel_az_pred,
-        'rel_el_pred'  : rel_el_pred,
-        'rel_ro_pred'  : rel_ro_pred,
-    }
-    
-    return ans_dict 
-
-# input PIL Image
-@spaces.GPU
-@torch.no_grad()
-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(batch=image_tensors.unsqueeze(0))
-    print(ans_dict)
-    return ans_dict