moved flax model and conversion code to separate repository

replicate/cog.yaml

@@ -0,0 +1,12 @@
+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

@@ -0,0 +1,23 @@
+import tempfile
+from cog import BasePredictor, Path, Input
+
+from min_dalle.min_dalle_torch import MinDalleTorch
+
+class Predictor(BasePredictor):
+    def setup(self):
+        self.model = MinDalleTorch(is_mega=True)
+
+    def predict(
+        self,
+        text: str = Input(
+            description="Text for generating images.",
+        ),
+        seed: int = Input(
+            description="Specify the seed.",
+        ),
+    ) -> Path:
+        image = self.model.generate_image(text, seed)
+        out_path = Path(tempfile.mkdtemp()) / "output.png"
+        image.save(str(out_path))
+
+        return out_path