Merge branch 'main' of github.com:pesser/stable-diffusion into main

configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512.yaml

@@ -2,7 +2,6 @@ model:
   base_learning_rate: 1.0e-04
   target: ldm.models.diffusion.ddpm.LatentDiffusion
   params:
-    ckpt_path: "/home/mchorse/stable-diffusion-ckpts/256pretrain-2022-06-09.ckpt"
     linear_start: 0.00085
     linear_end: 0.0120
     num_timesteps_cond: 1
@@ -20,7 +19,7 @@ model:
     scheduler_config: # 10000 warmup steps
       target: ldm.lr_scheduler.LambdaLinearScheduler
       params:
-        warm_up_steps: [ 10000 ]
+        warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch
         cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
         f_start: [ 1.e-6 ]
         f_max: [ 1. ]

configs/stable-diffusion/v1_improvedaesthetics.yaml

@@ -0,0 +1,135 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
+
+
+data:
+  target: ldm.data.laion.WebDataModuleFromConfig
+  params:
+    tar_base: "__improvedaesthetic__"
+    batch_size: 4
+    num_workers: 4
+    multinode: True
+    train:
+      shards: '{00000..17279}.tar -'
+      shuffle: 10000
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 512
+          interpolation: 3
+      - target: torchvision.transforms.RandomCrop
+        params:
+          size: 512
+
+    # NOTE use enough shards to avoid empty validation loops in workers
+    validation:
+      shards: '{17280..17535}.tar -'
+      shuffle: 0
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 512
+          interpolation: 3
+      - target: torchvision.transforms.CenterCrop
+        params:
+          size: 512
+
+
+lightning:
+  find_unused_parameters: False
+
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 5000
+
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 5000
+        max_images: 4
+        increase_log_steps: False
+        log_first_step: False
+        log_images_kwargs:
+          use_ema_scope: False
+          inpaint: False
+          plot_progressive_rows: False
+          plot_diffusion_rows: False
+          N: 4
+          unconditional_guidance_scale: 3.0
+          unconditional_guidance_label: [""]
+
+  trainer:
+    benchmark: True
+    val_check_interval: 5000000  # really sorry
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 2

configs/stable-diffusion/v2_laionhr1024.yaml

@@ -0,0 +1,132 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.001
+    linear_end: 0.015
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 16
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.22765929   # magic number
+  
+    # NOTE disabled for resuming
+    #scheduler_config: # 10000 warmup steps
+    #  target: ldm.lr_scheduler.LambdaLinearScheduler
+    #  params:
+    #    warm_up_steps: [ 10000 ]
+    #    cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+    #    f_start: [ 1.e-6 ]
+    #    f_max: [ 1. ]
+    #    f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 64    # not really needed
+        in_channels: 16
+        out_channels: 16
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 16
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: True
+          z_channels: 16
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [ 1,1,2,2,4 ]  # num_down = len(ch_mult)-1
+          num_res_blocks: 2
+          attn_resolutions: [ 16 ]
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
+
+
+data:
+  target: ldm.data.laion.WebDataModuleFromConfig
+  params:
+    tar_base: "pipe:aws s3 cp s3://s-datasets/laion-high-resolution/"
+    batch_size: 3
+    num_workers: 4
+    multinode: True
+    train:
+      shards: '{00000..17279}.tar -'
+      shuffle: 10000
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 1024
+          interpolation: 3
+      - target: torchvision.transforms.RandomCrop
+        params:
+          size: 1024
+
+    # NOTE use enough shards to avoid empty validation loops in workers
+    validation:
+      shards: '{17280..17535}.tar -'
+      shuffle: 0
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 1024
+          interpolation: 3
+      - target: torchvision.transforms.CenterCrop
+        params:
+          size: 1024
+
+
+lightning:
+  find_unused_parameters: False
+
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 2000
+
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 2000
+        max_images: 2
+        increase_log_steps: False
+        log_first_step: False
+        log_images_kwargs:
+          use_ema_scope: False
+          inpaint: False
+          plot_progressive_rows: False
+          plot_diffusion_rows: False
+          N: 2
+          unconditional_guidance_scale: 5.0
+          unconditional_guidance_label: [""]
+
+  trainer:
+    benchmark: True
+    val_check_interval: 5000000
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 4

configs/stable-diffusion/v3_pretraining.yaml

@@ -0,0 +1,137 @@
+model:
+  base_learning_rate: 8.e-05
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 32
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 416
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: [ 2, 2, 2, 2 ]
+        channel_mult: [ 1, 2, 4, 4 ]
+        disable_self_attentions: [ False, False, False, False ]  # converts the self-attention to a cross-attention layer if true
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ckpt_path: "/fsx/stable-diffusion/stable-diffusion/models/first_stage_models/kl-f8/model.ckpt"
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
+
+
+data:
+  target: ldm.data.laion.WebDataModuleFromConfig
+  params:
+    tar_base: "__improvedaesthetic__"
+    batch_size: 8
+    num_workers: 4
+    multinode: True
+    train:
+      shards: '{00000..17279}.tar -'
+      shuffle: 10000
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 256
+          interpolation: 3
+      - target: torchvision.transforms.RandomCrop
+        params:
+          size: 256
+
+#    # NOTE use enough shards to avoid empty validation loops in workers
+    validation:
+      shards: '{17280..17535}.tar -'
+      shuffle: 0
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 256
+          interpolation: 3
+      - target: torchvision.transforms.CenterCrop
+        params:
+          size: 256
+
+
+lightning:
+  find_unused_parameters: false
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 5000
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: True
+        batch_frequency: 2500  
+        max_images: 4
+        increase_log_steps: False
+        log_first_step: False
+        log_images_kwargs:
+          use_ema_scope: False
+          inpaint: False
+          plot_progressive_rows: False
+          plot_diffusion_rows: False
+          N: 4
+          unconditional_guidance_scale: 3.0
+          unconditional_guidance_label: [""]
+
+  trainer:
+    #replace_sampler_ddp: False
+    benchmark: True
+    val_check_interval: 5000000 # really sorry
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1

ldm/data/laion.py

@@ -151,12 +151,7 @@ class WebDataModuleFromConfig(pl.LightningDataModule):
         if self.tar_base == "__improvedaesthetic__":
             print("## Warning, loading the same improved aesthetic dataset "
                     "for all splits and ignoring shards parameter.")
-            urls = []
-            for i in range(1, 65):
-                for j in range(512):
-                    for k in range(5):
-                        urls.append(f's3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics/{i:02d}/{j:03d}/{k:05d}.tar')
-            tars = [f'pipe:aws s3 cp {url} -' for url in urls]
+            tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
         else:
             tars = os.path.join(self.tar_base, dataset_config.shards)
 
@@ -314,7 +309,8 @@ if __name__ == "__main__":
     from pytorch_lightning.trainer.supporters import CombinedLoader, CycleIterator
 
     #config = OmegaConf.load("configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512.yaml")
-    config = OmegaConf.load("configs/stable-diffusion/txt2img-upscale-clip-encoder-f16-1024.yaml")
+    #config = OmegaConf.load("configs/stable-diffusion/txt2img-upscale-clip-encoder-f16-1024.yaml")
+    config = OmegaConf.load("configs/stable-diffusion/txt2img-v2-clip-encoder-improved_aesthetics-256.yaml")
     datamod = WebDataModuleFromConfig(**config["data"]["params"])
     dataloader = datamod.train_dataloader()
 

main.py

@@ -21,6 +21,9 @@ from ldm.data.base import Txt2ImgIterableBaseDataset
 from ldm.util import instantiate_from_config
 
 
+MULTINODE_HACKS = True
+
+
 def get_parser(**parser_kwargs):
     def str2bool(v):
         if isinstance(v, bool):
@@ -268,6 +271,9 @@ class SetupCallback(Callback):
                     os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True)
             print("Project config")
             print(OmegaConf.to_yaml(self.config))
+            if MULTINODE_HACKS:
+                import time
+                time.sleep(5)
             OmegaConf.save(self.config,
                            os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)))
 
@@ -278,7 +284,7 @@ class SetupCallback(Callback):
 
         else:
             # ModelCheckpoint callback created log directory --- remove it
-            if not self.resume and os.path.exists(self.logdir):
+            if not MULTINODE_HACKS and not self.resume and os.path.exists(self.logdir):
                 dst, name = os.path.split(self.logdir)
                 dst = os.path.join(dst, "child_runs", name)
                 os.makedirs(os.path.split(dst)[0], exist_ok=True)
@@ -759,9 +765,19 @@ if __name__ == "__main__":
             del callbacks_cfg['ignore_keys_callback']
 
         trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg]
+        if not "plugins" in trainer_kwargs:
+            trainer_kwargs["plugins"] = list()
         if not lightning_config.get("find_unused_parameters", True):
             from pytorch_lightning.plugins import DDPPlugin
-            trainer_kwargs["plugins"] = DDPPlugin(find_unused_parameters=False)
+            trainer_kwargs["plugins"].append(DDPPlugin(find_unused_parameters=False))
+        if MULTINODE_HACKS:
+            # disable resume from hpc ckpts
+            # NOTE below only works in later versions
+            # from pytorch_lightning.plugins.environments import SLURMEnvironment
+            # trainer_kwargs["plugins"].append(SLURMEnvironment(auto_requeue=False))
+            # hence we monkey patch things
+            from pytorch_lightning.trainer.connectors.checkpoint_connector import CheckpointConnector
+            setattr(CheckpointConnector, "hpc_resume_path", None)
 
         trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
         trainer.logdir = logdir  ###
@@ -833,6 +849,18 @@ if __name__ == "__main__":
                 raise
         if not opt.no_test and not trainer.interrupted:
             trainer.test(model, data)
+    except RuntimeError as err:
+        if MULTINODE_HACKS:
+            import requests
+            import datetime
+            import os
+            import socket
+            device = os.environ.get("CUDA_VISIBLE_DEVICES", "?")
+            hostname = socket.gethostname()
+            ts = datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')
+            resp = requests.get('http://169.254.169.254/latest/meta-data/instance-id')
+            print(f'ERROR at {ts} on {hostname}/{resp.text} (CUDA_VISIBLE_DEVICES={device}): {type(err).__name__}: {err}', flush=True)
+        raise err
     except Exception:
         if opt.debug and trainer.global_rank == 0:
             try:

scripts/checker.py

@@ -0,0 +1,218 @@
+import os
+import glob
+import subprocess
+import time
+import fire
+
+import numpy as np
+from tqdm import tqdm
+import torch
+from pytorch_lightning import seed_everything
+from omegaconf import OmegaConf
+from ldm.util import instantiate_from_config
+from ldm.models.diffusion.plms import PLMSSampler
+from einops import rearrange
+from torchvision.utils import make_grid
+from PIL import Image
+import contextlib
+
+
+def load_model_from_config(config, ckpt, verbose=False):
+    pl_sd = torch.load(ckpt, map_location="cpu")
+    gs = pl_sd["global_step"]
+    sd = pl_sd["state_dict"]
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=True)
+    model.cuda()
+    model.eval()
+    return model, gs
+
+
+def read_prompts(path):
+    with open(path, "r") as f:
+        prompts = f.read().splitlines()
+    return prompts
+
+
+def split_in_batches(iterator, n):
+    out = []
+    for elem in iterator:
+        out.append(elem)
+        if len(out) == n:
+            yield out
+            out = []
+    if len(out) > 0:
+        yield out
+
+
+class Sampler(object):
+    def __init__(self, out_dir, ckpt_path, cfg_path, prompts_path, shape, seed=42):
+        self.out_dir = out_dir
+        self.ckpt_path = ckpt_path
+        self.cfg_path = cfg_path
+        self.prompts_path = prompts_path
+        self.seed = seed
+
+        self.batch_size = 1
+        self.scale = 10
+        self.shape = shape
+        self.n_steps = 100
+        self.nrow = 8
+
+
+    @torch.inference_mode()
+    def sample(self, model, prompts, ema=True):
+        seed = self.seed
+        batch_size = self.batch_size
+        scale = self.scale
+        n_steps = self.n_steps
+
+        shape = self.shape
+
+        print("Sampling model.")
+        print("ckpt_path", self.ckpt_path)
+        print("cfg_path", self.cfg_path)
+        print("prompts_path", self.prompts_path)
+        print("out_dir", self.out_dir)
+        print("seed", self.seed)
+        print("batch_size", batch_size)
+        print("scale", scale)
+        print("n_steps", n_steps)
+        print("shape", shape)
+
+        prompts = list(split_in_batches(prompts, batch_size))
+
+        sampler = PLMSSampler(model)
+        all_samples = list()
+
+        ctxt = model.ema_scope if ema else contextlib.nullcontext
+
+        with ctxt():
+            for prompts_batch in tqdm(prompts, desc="prompts"):
+                uc = None
+                if scale != 1.0:
+                    uc = model.get_learned_conditioning(batch_size * [""])
+                c = model.get_learned_conditioning(prompts_batch)
+
+                seed_everything(seed)
+
+                samples_latent, _ = sampler.sample(
+                    S=n_steps,
+                    conditioning=c,
+                    batch_size=batch_size,
+                    shape=shape,
+                    verbose=False,
+                    unconditional_guidance_scale=scale,
+                    unconditional_conditioning=uc,
+                    eta=0.0,
+                    dynamic_threshold=None,
+                )
+
+                samples = model.decode_first_stage(samples_latent)
+                samples = torch.clamp((samples+1.0)/2.0, min=0.0, max=1.0)
+
+                all_samples.append(samples)
+
+        all_samples = torch.cat(all_samples, 0)
+        return all_samples
+
+
+    @torch.inference_mode()
+    def __call__(self):
+        config = OmegaConf.load(self.cfg_path)
+        model, global_step = load_model_from_config(config, self.ckpt_path)
+        print(f"Restored model at global step {global_step}.")
+
+        prompts = read_prompts(self.prompts_path)
+
+        all_samples = self.sample(model, prompts, ema=True)
+        self.save_as_grid("grid_with_wings", all_samples, global_step)
+        all_samples = self.sample(model, prompts, ema=False)
+        self.save_as_grid("grid_without_wings", all_samples, global_step)
+
+
+    def save_as_grid(self, name, grid, global_step):
+        grid = make_grid(grid, nrow=self.nrow)
+        grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
+
+        os.makedirs(self.out_dir, exist_ok=True)
+        filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(
+            name,
+            global_step,
+            0,
+            0,
+        )
+        grid_path = os.path.join(self.out_dir, filename)
+        Image.fromarray(grid.astype(np.uint8)).save(grid_path)
+        print(f"---> {grid_path}")
+
+
+class Checker(object):
+    def __init__(self, ckpt_path, callback, wait_for_file=5, interval=60):
+        self._cached_stamp = 0
+        self.filename = ckpt_path
+        self.callback = callback
+        self.interval = interval
+        self.wait_for_file = wait_for_file
+
+    def check(self):
+        while True:
+            stamp = os.stat(self.filename).st_mtime
+            if stamp != self._cached_stamp:
+                while True:
+                    # try to wait until checkpoint is fully written
+                    previous_stamp = stamp
+                    time.sleep(self.wait_for_file)
+                    stamp = os.stat(self.filename).st_mtime
+                    if stamp != previous_stamp:
+                        print(f"File is still changing. Waiting {self.wait_for_file} seconds.")
+                    else:
+                        break
+
+                self._cached_stamp = stamp
+                # file has changed, so do something...
+                print(f"{self.__class__.__name__}: Detected a new file at "
+                      f"{self.filename}, calling back.")
+                self.callback()
+
+            else:
+                time.sleep(self.interval)
+
+
+def run(prompts_path="scripts/prompts/prompts-with-wings.txt",
+        watch_log_dir=None, out_dir=None, ckpt_path=None, cfg_path=None,
+        H=256,
+        W=None,
+        C=4,
+        F=8,
+        wait_for_file=5,
+        interval=60):
+
+    if out_dir is None:
+        assert watch_log_dir is not None
+        out_dir = os.path.join(watch_log_dir, "images/checker")
+
+    if ckpt_path is None:
+        assert watch_log_dir is not None
+        ckpt_path = os.path.join(watch_log_dir, "checkpoints/last.ckpt")
+
+    if cfg_path is None:
+        assert watch_log_dir is not None
+        configs = glob.glob(os.path.join(watch_log_dir, "configs/*-project.yaml"))
+        cfg_path = sorted(configs)[-1]
+
+    if W is None:
+        assert H is not None
+        W = H
+    if H is None:
+        assert W is not None
+        H = W
+    shape = [C, H//F, W//F]
+    sampler = Sampler(out_dir, ckpt_path, cfg_path, prompts_path, shape=shape)
+
+    checker = Checker(ckpt_path, sampler, wait_for_file=wait_for_file, interval=interval)
+    checker.check()
+
+
+if __name__ == "__main__":
+    fire.Fire(run)

scripts/slurm/v1_improvedaesthetics/launcher.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+
+# mpi version for node rank
+H=`hostname`
+THEID=`echo -e $HOSTNAMES  | python3 -c "import sys;[sys.stdout.write(str(i)) for i,line in enumerate(next(sys.stdin).split(' ')) if line.strip() == '$H'.strip()]"`
+export NODE_RANK=${THEID}
+echo THEID=$THEID
+
+echo "##########################################"
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo NODE_RANK=${NODE_RANK}
+echo WORLD_SIZE=${WORLD_SIZE}
+echo "##########################################"
+# debug environment worked great so we stick with it
+# no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
+# env with pip dependencies from stable diffusion's requirements.txt
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+conda activate stable
+cd /fsx/stable-diffusion/stable-diffusion
+
+CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v1_improvedaesthetics.yaml"
+
+# resume and set new seed to reshuffle data
+EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-11T20-16-11_txt2img-1p4B-multinode-clip-encoder-high-res-512_improvedaesthetic/checkpoints/last.ckpt"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v1_improvedaesthetics/sbatch.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v1-improvedaesthetics
+#SBATCH --nodes 20
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-gpu=4
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --output=%x_%j.out
+#SBATCH --comment "Key=Monitoring,Value=ON"
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+export NCCL_DEBUG=info
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+export WORLD_SIZE=$COUNT_NODE
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+echo $WORLD_SIZE
+
+mpirun -n $COUNT_NODE -perhost 1 /fsx/stable-diffusion/stable-diffusion/scripts/slurm/v1_improvedaesthetics/launcher.sh

scripts/slurm/v2_laionhr1024/launcher.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+
+# mpi version for node rank
+H=`hostname`
+THEID=`echo -e $HOSTNAMES  | python3 -c "import sys;[sys.stdout.write(str(i)) for i,line in enumerate(next(sys.stdin).split(' ')) if line.strip() == '$H'.strip()]"`
+export NODE_RANK=${THEID}
+echo THEID=$THEID
+
+echo "##########################################"
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo NODE_RANK=${NODE_RANK}
+echo WORLD_SIZE=${WORLD_SIZE}
+echo "##########################################"
+# debug environment worked great so we stick with it
+# no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
+# env with pip dependencies from stable diffusion's requirements.txt
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+conda activate stable
+cd /fsx/stable-diffusion/stable-diffusion
+
+CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v2_laionhr1024.yaml"
+
+# resume and set new seed to reshuffle data
+EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-12T00-50-44_txt2img-multinode-clip-encoder-f16-1024-laion-hr/checkpoints/last.ckpt"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v2_laionhr1024/sbatch.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v2-laionhr1024
+#SBATCH --nodes 20
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-gpu=4
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --output=%x_%j.out
+#SBATCH --comment "Key=Monitoring,Value=ON"
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+export NCCL_DEBUG=info
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+export WORLD_SIZE=$COUNT_NODE
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+echo $WORLD_SIZE
+
+mpirun -n $COUNT_NODE -perhost 1 /fsx/stable-diffusion/stable-diffusion/scripts/slurm/v2_laionhr1024/launcher.sh

scripts/slurm/v3_pretraining/launcher.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+
+# mpi version for node rank
+H=`hostname`
+THEID=`echo -e $HOSTNAMES  | python3 -c "import sys;[sys.stdout.write(str(i)) for i,line in enumerate(next(sys.stdin).split(' ')) if line.strip() == '$H'.strip()]"`
+export NODE_RANK=${THEID}
+echo THEID=$THEID
+
+echo "##########################################"
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo NODE_RANK=${NODE_RANK}
+echo WORLD_SIZE=${WORLD_SIZE}
+echo "##########################################"
+# debug environment worked great so we stick with it
+# no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
+# env with pip dependencies from stable diffusion's requirements.txt
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+conda activate stable
+cd /fsx/stable-diffusion/stable-diffusion
+
+CONFIG=configs/stable-diffusion/v3_pretraining.yaml
+
+# resume and set new seed to reshuffle data
+EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/rlogs/2022-07-11T22-57-10_txt2img-v2-clip-encoder-improved_aesthetics-256/checkpoints/last.ckpt"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v3_pretraining/sbatch.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v3-pretraining
+#SBATCH --nodes 44
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-gpu=4
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --output=%x_%j.out
+#SBATCH --comment "Key=Monitoring,Value=ON"
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+export NCCL_DEBUG=info
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+export WORLD_SIZE=$COUNT_NODE
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+echo $WORLD_SIZE
+
+mpirun -n $COUNT_NODE -perhost 1 /fsx/stable-diffusion/stable-diffusion/scripts/slurm/v3_pretraining/launcher.sh

scripts/txt2img.py

@@ -40,7 +40,7 @@ def load_model_from_config(config, ckpt, verbose=False):
     return model
 
 
-if __name__ == "__main__":
+def main():
     parser = argparse.ArgumentParser()
 
     parser.add_argument(
@@ -258,3 +258,7 @@ if __name__ == "__main__":
     print(f"Your samples are ready and waiting for you here: \n{outpath} \n"
           f"Sampling took {toc-tic}s, i.e. produced {opt.n_iter * opt.n_samples / (toc - tic):.2f} samples/sec."
           f" \nEnjoy.")
+
+
+if __name__ == "__main__":
+    main()