diff --git a/big_sleep/big_sleep.py b/big_sleep/big_sleep.py index 68052d6..ffc4b06 100644 --- a/big_sleep/big_sleep.py +++ b/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: diff --git a/big_sleep/cli.py b/big_sleep/cli.py index 81c2eeb..8253836 100644 --- a/big_sleep/cli.py +++ b/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(): diff --git a/big_sleep/version.py b/big_sleep/version.py index 13a85f7..44b1806 100644 --- a/big_sleep/version.py +++ b/big_sleep/version.py @@ -1 +1 @@ -__version__ = '0.2.5' +__version__ = '0.2.6'