Update ppd/models/ppd.py

ppd/models/ppd.py

@@ -53,12 +53,20 @@ class PixelPerfectDepth(nn.Module):
         )
     
     @torch.no_grad()
-    def forward_test(self, image):
+    def infer_image(self, image):
         h, w = image.shape[:2]
         image = resize_keep_aspect(image)
         image = image2tensor(image)
         image = image.to(self.device)
 
+        depth = self.forward_test(image)
+        depth = F.interpolate(depth, size=(h, w), mode='bilinear', align_corners=False)[0, 0]
+
+        return depth.cpu().numpy()
+    
+    @torch.no_grad()
+    def forward_test(self, image):
+
         semantics = self.semantics_prompt(image)
         cond = image - 0.5
         latent = torch.randn(size=[cond.shape[0], 1, cond.shape[2], cond.shape[3]]).to(self.device)
@@ -68,10 +76,7 @@ class PixelPerfectDepth(nn.Module):
             pred = self.dit(x=input, semantics=semantics, timestep=timestep)
             latent = self.sampler.step(pred=pred, x_t=latent, t=timestep)
 
-        depth = latent + 0.5
-        depth = F.interpolate(depth, size=(h, w), mode='bilinear', align_corners=False)[0, 0]
-
-        return depth.cpu().numpy()
+        return latent + 0.5
 
 
     @torch.no_grad()