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

ldm/models/diffusion/ddpm.py

@@ -74,6 +74,7 @@ class DDPM(pl.LightningModule):
                  learn_logvar=False,
                  logvar_init=0.,
                  make_it_fit=False,
+                 ucg_training=None,
                  ):
         super().__init__()
         assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
@@ -117,6 +118,10 @@ class DDPM(pl.LightningModule):
         if self.learn_logvar:
             self.logvar = nn.Parameter(self.logvar, requires_grad=True)
 
+        self.ucg_training = ucg_training or dict()
+        if self.ucg_training:
+            self.ucg_prng = np.random.RandomState()
+
 
     def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
                           linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
@@ -389,6 +394,15 @@ class DDPM(pl.LightningModule):
         return loss, loss_dict
 
     def training_step(self, batch, batch_idx):
+        for k in self.ucg_training:
+            p = self.ucg_training[k]["p"]
+            val = self.ucg_training[k]["val"]
+            if val is None:
+                val = ""
+            for i in range(len(batch[k])):
+                if self.ucg_prng.choice(2, p=[1-p, p]):
+                    batch[k][i] = val
+
         loss, loss_dict = self.shared_step(batch)
 
         self.log_dict(loss_dict, prog_bar=True,

scripts/slurm/v1_iahr_torch111/launcher.sh

@@ -26,7 +26,8 @@ CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v1_impro
 #EXTRA="--seed 718 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/checkpoints2/v1pp/v1pp-flatline.ckpt"
 #EXTRA="--seed 718 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-22T07-45-07_v1_improvedaesthetics/checkpoints/last.ckpt"
 #EXTRA="--seed 719 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-22T12-32-32_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
-EXTRA="--seed 720 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-23T07-52-21_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
+#EXTRA="--seed 720 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-23T07-52-21_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
+EXTRA="--seed 721 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-24T19-07-33_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
 
 # only images >= 512 and pwatermark <= 0.4999
 EXTRA="${EXTRA} data.params.min_size=512 data.params.max_pwatermark=0.4999"
@@ -35,7 +36,7 @@ EXTRA="${EXTRA} data.params.min_size=512 data.params.max_pwatermark=0.4999"
 EXTRA="${EXTRA} -f v1_iahr_torch111"
 
 # time to decay
-#EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[50000] model.params.scheduler_config.params.f_min=[1e-6]"
+EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[300000] model.params.scheduler_config.params.warm_up_steps=[250000] model.params.scheduler_config.params.f_min=[1e-6]"
 
 # custom logdir
 #EXTRA="${EXTRA} --logdir rlogs"

scripts/slurm/v1_iahr_torch111_ucg/launcher.sh

@@ -0,0 +1,52 @@
+#!/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
+conda activate torch111
+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 718 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/checkpoints2/v1pp/v1pp-flatline.ckpt"
+#EXTRA="--seed 718 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-22T07-45-07_v1_improvedaesthetics/checkpoints/last.ckpt"
+#EXTRA="--seed 719 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-22T12-32-32_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
+#EXTRA="--seed 720 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-23T07-52-21_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
+EXTRA="--seed 721 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-24T19-07-33_v1_improvedaestheticsv1_iahr_torch111/checkpoints/last.ckpt"
+
+# only images >= 512 and pwatermark <= 0.4999
+EXTRA="${EXTRA} data.params.min_size=512 data.params.max_pwatermark=0.4999"
+
+# unconditional guidance training
+EXTRA="${EXTRA} model.params.ucg_training.txt.p=0.1 model.params.ucg_training.txt.val=''"
+
+# postfix
+EXTRA="${EXTRA} -f v1_iahr_torch111_ucg"
+
+# time to decay
+#EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[300000] model.params.scheduler_config.params.warm_up_steps=[250000] model.params.scheduler_config.params.f_min=[1e-6]"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+/bin/bash /fsx/stable-diffusion/stable-diffusion/scripts/test_gpu.sh
+
+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_iahr_torch111_ucg/sbatch.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v1-iahr-torch111-ucg
+#SBATCH --nodes 32
+#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"
+#SBATCH --no-requeue
+
+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-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+#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_iahr_torch111_ucg/launcher.sh