update replicate files

replicate/cog.yaml

@@ -1,12 +0,0 @@
-build:
-  cuda: "11.0"
-  gpu: true
-  python_version: "3.8"
-  system_packages:
-    - "libgl1-mesa-glx"
-    - "libglib2.0-0"
-  python_packages:
-    - "torch==1.10.1"
-    - "flax==0.5.2"
-
-predict: "predict.py:Predictor"

replicate/predict.py

@@ -15,8 +15,11 @@ class Predictor(BasePredictor):
         seed: int = Input(
             description="Specify a random seed.",
         ),
+        grid_size: int = Input(
+            description="Specify the grid size.",
+        )
     ) -> Path:
-        image = self.model.generate_image(text, seed, grid_size=3)
+        image = self.model.generate_image(text, seed, grid_size=grid_size)
         out_path = Path(tempfile.mkdtemp()) / "output.png"
         image.save(str(out_path))