moved flax model and conversion code to separate repository

This commit is contained in:
Brett Kuprel
2022-07-01 14:06:50 -04:00
parent febd18df77
commit 07ce93d5f8
13 changed files with 57 additions and 712 deletions
+23
View File
@@ -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