From 9cde5b1a6050ed567785f971866ef954adb7c872 Mon Sep 17 00:00:00 2001 From: Brett Kuprel Date: Sat, 2 Jul 2022 10:32:17 -0400 Subject: [PATCH] update replicate files --- cog.yaml | 2 +- replicate/predict.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/cog.yaml b/cog.yaml index 42e8f9b..d835b67 100644 --- a/cog.yaml +++ b/cog.yaml @@ -9,4 +9,4 @@ build: - "torch==1.10.1" - "flax==0.5.2" -predict: "predict.py:Predictor" +predict: "replicate/predict.py:Predictor" diff --git a/replicate/predict.py b/replicate/predict.py index 794d563..395d6f6 100644 --- a/replicate/predict.py +++ b/replicate/predict.py @@ -13,10 +13,12 @@ class Predictor(BasePredictor): description="Text for generating images.", ), seed: int = Input( - description="Specify a random seed.", + description="Specify a random seed." ), grid_size: int = Input( description="Specify the grid size.", + ge=1, + le=4 ) ) -> Path: image = self.model.generate_image(text, seed, grid_size=grid_size)