auto open folder where imaginations will be saved to

big_sleep/big_sleep.py

@@ -3,16 +3,20 @@
 from torch import nn
 from torch.optim import Adam
 
-from pathlib import Path
-from tqdm import trange
 import torchvision
 from torchvision.utils import save_image
 
+import os
+import sys
+import subprocess
+import signal
+from pathlib import Path
+from tqdm import trange
+from collections import namedtuple
+
 from big_sleep.biggan import BigGAN
 from big_sleep.clip import load, tokenize, normalize_image
 
-import signal
-from collections import namedtuple
 from einops import rearrange
 
 assert torch.cuda.is_available(), 'CUDA must be available in order to use Deep Daze'
@@ -27,6 +31,32 @@ def signal_handling(signum,frame):
 
 signal.signal(signal.SIGINT,signal_handling)
 
+# helpers
+
+def open_folder(path):
+    if os.path.isfile(path):
+        path = os.path.dirname(path)
+
+    if not os.path.isdir(path):
+        return
+
+    cmd_list = None
+    if sys.platform == 'darwin':
+        cmd_list = ['open', '--', path]
+    elif sys.platform == 'linux2' or sys.platform == 'linux':
+        cmd_list = ['xdg-open', path]
+    elif sys.platform in ['win32', 'win64']:
+        cmd_list = ['explorer', path.replace('/','\\')]
+    if cmd_list == None:
+        return
+
+    try:
+        subprocess.check_call(cmd_list)
+    except subprocess.CalledProcessError:
+        pass
+    except OSError:
+        pass
+
 # load clip
 
 perceptor, preprocess = load()
@@ -134,7 +164,8 @@ def __init__(
         epochs = 20,
         iterations = 1050,
         save_progress = False,
-        bilinear = False
+        bilinear = False,
+        open_folder = True
     ):
         super().__init__()
         self.epochs = epochs
@@ -159,6 +190,7 @@ def __init__(
         self.save_progress = save_progress
 
         self.encoded_text = tokenize(text).cuda()
+        self.open_folder = open_folder
 
     def train_step(self, epoch, i):
         total_loss = 0
@@ -188,6 +220,9 @@ def train_step(self, epoch, i):
     def forward(self):
         print(f'Imagining "{self.text}" from the depths of my weights...')
 
+        if self.open_folder:
+            open_folder('./')
+
         for epoch in trange(self.epochs, desc = 'epochs'):
             pbar = trange(self.iterations, desc='iteration')
             for i in pbar:

big_sleep/cli.py

@@ -12,7 +12,8 @@ def train(
     save_every = 50,
     overwrite = False,
     save_progress = False,
-    bilinear = False
+    bilinear = False,
+    open_folder = True
 ):
 
     imagine = Imagine(
@@ -24,7 +25,8 @@ def train(
         iterations = iterations,
         save_every = save_every,
         save_progress = save_progress,
-        bilinear = bilinear
+        bilinear = bilinear,
+        open_folder = open_folder
     )
 
     if not overwrite and imagine.filename.exists():

big_sleep/version.py

@@ -1 +1 @@
-__version__ = '0.2.5'
+__version__ = '0.2.6'