github.com/sieve-data/cog@v0.0.72/pkg/cli/init-templates/predict.py (about) 1 # Prediction interface for Cog ⚙️ 2 # https://github.com/sieve-data/cog/blob/main/docs/python.md 3 4 from cog import BasePredictor, Input, Path 5 6 7 class Predictor(BasePredictor): 8 def setup(self): 9 """Load the model into memory to make running multiple predictions efficient""" 10 # self.model = torch.load("./weights.pth") 11 12 def predict( 13 self, 14 image: Path = Input(description="Grayscale input image"), 15 scale: float = Input( 16 description="Factor to scale image by", ge=0, le=10, default=1.5 17 ), 18 ) -> Path: 19 """Run a single prediction on the model""" 20 # processed_input = preprocess(image) 21 # output = self.model(processed_image, scale) 22 # return postprocess(output)