diff --git a/replicate/cog.yaml b/cog.yaml
similarity index 100%
rename from replicate/cog.yaml
rename to cog.yaml
diff --git a/replicate/predict.py b/replicate/predict.py
index 7be6f94..794d563 100644
--- a/replicate/predict.py
+++ b/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))
 
diff --git a/setup.py b/setup.py
index aced9b4..ec38227 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ setuptools.setup(
     name='min-dalle',
     description = 'min(DALL·E)',
     long_description=(Path(__file__).parent / "README").read_text(),
-    version='0.2.9',
+    version='0.2.10',
     author='Brett Kuprel',
     author_email='brkuprel@gmail.com',
     url='https://github.com/kuprel/min-dalle',