generate_images_stream and generate_images

2022-07-07 17:03:47 -04:00
parent b17bea11b6
commit 2cac9220b5
5 changed files with 62 additions and 111 deletions
@@ -1,5 +1,6 @@
 from min_dalle import MinDalle
 import tempfile
+import torch
 from typing import Iterator
 from cog import BasePredictor, Path, Input

@@ -53,5 +54,6 @@ class ReplicatePredictor(BasePredictor):
        except:
            print("An error occured, deleting model")
            del self.model
+            torch.cuda.empty_cache()
            self.setup()
            raise Exception("There was an error, please try again")