update cog init predictor template

Signed-off-by: Zeke Sikelianos <[email protected]>

Author: zeke

pkg/cli/init-templates/predict.py

@@ -1,18 +1,22 @@
 # Prediction interface for Cog ⚙️
-# Reference: https://github.com/replicate/cog/blob/main/docs/python.md
+# https://github.com/replicate/cog/blob/main/docs/python.md
 
-import cog
-# import torch
+from cog import BasePredictor, Input, Path
 
-class Predictor(cog.Predictor):
+
+class Predictor(BasePredictor):
     def setup(self):
-      """Load the model into memory to make running multiple predictions efficient"""
-      # self.model = torch.load("./weights.pth")
+        """Load the model into memory to make running multiple predictions efficient"""
+        # self.model = torch.load("./weights.pth")
 
-    @cog.input("image", type=cog.Path, help="Grayscale input image")
-    @cog.input("scale", type=float, default=1.5, help="Factor to scale image by")
-    def predict(self, image):
+    def predict(
+        self,
+        input: Path = Input(title="Grayscale input image"),
+        scale: float = Input(
+            title="Factor to scale image by", gt=0, lt=10, default=1.5
+        ),
+    ) -> Path:
         """Run a single prediction on the model"""
-        # processed_input = preprocess(image)
-        # output = self.model(processed_input)
-        # return post_processing(output)
+        # processed_input = preprocess(input)
+        # output = self.model(processed_input, scale)
+        # return postprocess(output)