From b17bea11b60fd14c5bfc4525936a8f245ca16057 Mon Sep 17 00:00:00 2001
From: Brett Kuprel <brkuprel@gmail.com>
Date: Thu, 7 Jul 2022 12:35:00 -0400
Subject: [PATCH] add try-except to replicate

---
 cog.yaml               |  2 +-
 replicate_predictor.py | 40 +++++++++++++++++++++++-----------------
 2 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/cog.yaml b/cog.yaml
index c0551d5..61f36df 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -1,5 +1,5 @@
 build:
-  cuda: "11.5.1"
+  cuda: "11.4"
   gpu: true
   python_version: "3.10"
   system_packages:
diff --git a/replicate_predictor.py b/replicate_predictor.py
index d9ab027..bc2efa3 100644
--- a/replicate_predictor.py
+++ b/replicate_predictor.py
@@ -31,21 +31,27 @@ class ReplicatePredictor(BasePredictor):
             default=4
         ),
     ) -> Iterator[Path]:
-        seed = -1
-        log2_mid_count = 3 if intermediate_outputs else 0
-        image_stream = self.model.generate_image_stream(
-            text,
-            seed,
-            grid_size=grid_size,
-            log2_mid_count=log2_mid_count,
-            log2_supercondition_factor=log2_supercondition_factor,
-            is_verbose=True
-        )
+        try: 
+            seed = -1
+            log2_mid_count = 3 if intermediate_outputs else 0
+            image_stream = self.model.generate_image_stream(
+                text,
+                seed,
+                grid_size=grid_size,
+                log2_mid_count=log2_mid_count,
+                log2_supercondition_factor=log2_supercondition_factor,
+                is_verbose=True
+            )
 
-        iter = 0
-        path = Path(tempfile.mkdtemp())
-        for image in image_stream:
-            iter += 1
-            image_path = path / 'min-dalle-iter-{}.jpg'.format(iter)
-            image.save(str(image_path))
-            yield image_path
\ No newline at end of file
+            iter = 0
+            path = Path(tempfile.mkdtemp())
+            for image in image_stream:
+                iter += 1
+                image_path = path / 'min-dalle-iter-{}.jpg'.format(iter)
+                image.save(str(image_path))
+                yield image_path
+        except:
+            print("An error occured, deleting model")
+            del self.model
+            self.setup()
+            raise Exception("There was an error, please try again")
\ No newline at end of file