[WIP] polish code style of stable diffusion example

examples/images/diffusion/configs/Teyvat/train_colossalai_teyvat.yaml

@@ -1,6 +1,6 @@
 model:
   base_learning_rate: 1.0e-4
-  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  #target: ldm.models.diffusion.ddpm.LatentDiffusion
   params:
     parameterization: "v"
     linear_start: 0.00085
@@ -20,7 +20,7 @@ model:
     use_ema: False
 
     scheduler_config: # 10000 warmup steps
-      target: ldm.lr_scheduler.LambdaLinearScheduler
+      #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
@@ -30,7 +30,7 @@ model:
 
 
     unet_config:
-      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      #target: ldm.modules.diffusionmodules.openaimodel.UNetModel
       params:
         use_checkpoint: True
         use_fp16: True
@@ -49,7 +49,7 @@ model:
         legacy: False
 
     first_stage_config:
-      target: ldm.models.autoencoder.AutoencoderKL
+      #target: ldm.models.autoencoder.AutoencoderKL
       params:
         embed_dim: 4
         monitor: val/rec_loss
@@ -73,13 +73,13 @@ model:
           target: torch.nn.Identity
 
     cond_stage_config:
-      target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+      #target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
       params:
         freeze: True
         layer: "penultimate"
 
 data:
-  target: main.DataModuleFromConfig
+  #target: main.DataModuleFromConfig
   params:
     batch_size: 16
     num_workers: 4
@@ -105,7 +105,7 @@ lightning:
     precision: 16
     auto_select_gpus: False
     strategy:
-      target: strategies.ColossalAIStrategy
+      #target: strategies.ColossalAIStrategy
       params:
         use_chunk: True
         enable_distributed_storage: True
@@ -120,7 +120,7 @@ lightning:
 
   logger_config:
     wandb:
-      target: loggers.WandbLogger
+      #target: loggers.WandbLogger
       params:
           name: nowname
           save_dir: "/tmp/diff_log/"

examples/images/diffusion/configs/train_colossalai.yaml

@@ -1,6 +1,6 @@
 model:
   base_learning_rate: 1.0e-4
-  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  #target: ldm.models.diffusion.ddpm.LatentDiffusion
   params:
     parameterization: "v"
     linear_start: 0.00085
@@ -19,7 +19,7 @@ model:
     use_ema: False # we set this to false because this is an inference only config
 
     scheduler_config: # 10000 warmup steps
-      target: ldm.lr_scheduler.LambdaLinearScheduler
+      #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
@@ -29,7 +29,7 @@ model:
 
 
     unet_config:
-      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      #target: ldm.modules.diffusionmodules.openaimodel.UNetModel
       params:
         use_checkpoint: True
         use_fp16: True
@@ -48,7 +48,7 @@ model:
         legacy: False
 
     first_stage_config:
-      target: ldm.models.autoencoder.AutoencoderKL
+      #target: ldm.models.autoencoder.AutoencoderKL
       params:
         embed_dim: 4
         monitor: val/rec_loss
@@ -69,16 +69,16 @@ model:
           attn_resolutions: []
           dropout: 0.0
         lossconfig:
-          target: torch.nn.Identity
+          #target: torch.nn.Identity
 
     cond_stage_config:
-      target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+      #target: #ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
       params:
         freeze: True
         layer: "penultimate"
 
 data:
-  target: main.DataModuleFromConfig
+  #target: #main.DataModuleFromConfig
   params:
     batch_size: 128
     wrap: False
@@ -88,20 +88,20 @@ data:
     train:
       target: ldm.data.base.Txt2ImgIterableBaseDataset
       params:
-        file_path: # YOUR DATASET_PATH
+        file_path: /data/scratch/diffuser/laion_part0/
         world_size: 1
         rank: 0
 
 lightning:
   trainer:
     accelerator: 'gpu'
-    devices: 8
+    devices: 2
     log_gpu_memory: all
     max_epochs: 2
     precision: 16
     auto_select_gpus: False
     strategy:
-      target: strategies.ColossalAIStrategy
+      #target: #strategies.ColossalAIStrategy
       params:
         use_chunk: True
         enable_distributed_storage: True
@@ -116,7 +116,7 @@ lightning:
 
   logger_config:
     wandb:
-      target: loggers.WandbLogger
+      #target: #loggers.WandbLogger
       params:
           name: nowname
           save_dir: "/tmp/diff_log/"

examples/images/diffusion/configs/train_colossalai_cifar10.yaml

@@ -1,6 +1,6 @@
 model:
   base_learning_rate: 1.0e-4
-  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  #target: ldm.models.diffusion.ddpm.LatentDiffusion
   params:
     parameterization: "v"
     linear_start: 0.00085
@@ -19,7 +19,7 @@ model:
     use_ema: False # we set this to false because this is an inference only config
 
     scheduler_config: # 10000 warmup steps
-      target: ldm.lr_scheduler.LambdaLinearScheduler
+      #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
@@ -29,7 +29,7 @@ model:
 
 
     unet_config:
-      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      #target: ldm.modules.diffusionmodules.openaimodel.UNetModel
       params:
         use_checkpoint: True
         use_fp16: True
@@ -48,7 +48,7 @@ model:
         legacy: False
 
     first_stage_config:
-      target: ldm.models.autoencoder.AutoencoderKL
+      #target: ldm.models.autoencoder.AutoencoderKL
       params:
         embed_dim: 4
         monitor: val/rec_loss
@@ -69,16 +69,16 @@ model:
           attn_resolutions: []
           dropout: 0.0
         lossconfig:
-          target: torch.nn.Identity
+          #target: torch.nn.Identity
 
     cond_stage_config:
-      target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+      #target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
       params:
         freeze: True
         layer: "penultimate"
 
 data:
-  target: main.DataModuleFromConfig
+  #target: main.DataModuleFromConfig
   params:
     batch_size: 4
     num_workers: 4
@@ -105,7 +105,7 @@ lightning:
     precision: 16
     auto_select_gpus: False
     strategy:
-      target: strategies.ColossalAIStrategy
+      #target: strategies.ColossalAIStrategy
       params:
         use_chunk: True
         enable_distributed_storage: True

examples/images/diffusion/configs/train_ddp.yaml

@@ -1,6 +1,6 @@
 model:
   base_learning_rate: 1.0e-4
-  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  #target: ldm.models.diffusion.ddpm.LatentDiffusion
   params:
     parameterization: "v"
     linear_start: 0.00085
@@ -29,7 +29,7 @@ model:
 
 
     unet_config:
-      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      #target: ldm.modules.diffusionmodules.openaimodel.UNetModel
       params:
         use_checkpoint: True
         use_fp16: True
@@ -48,7 +48,7 @@ model:
         legacy: False
 
     first_stage_config:
-      target: ldm.models.autoencoder.AutoencoderKL
+      #target: ldm.models.autoencoder.AutoencoderKL
       params:
         embed_dim: 4
         monitor: val/rec_loss
@@ -72,13 +72,13 @@ model:
           target: torch.nn.Identity
 
     cond_stage_config:
-      target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+      #target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
       params:
         freeze: True
         layer: "penultimate"
 
 data:
-  target: main.DataModuleFromConfig
+  #target: main.DataModuleFromConfig
   params:
     batch_size: 128
     # num_workwers should be 2 * batch_size, and the total num less than 1024
@@ -100,7 +100,7 @@ lightning:
     precision: 16
     auto_select_gpus: False
     strategy:
-      target: strategies.DDPStrategy
+      #target: strategies.DDPStrategy
       params:
         find_unused_parameters: False
     log_every_n_steps: 2
@@ -111,7 +111,7 @@ lightning:
 
   logger_config:
     wandb:
-      target: loggers.WandbLogger
+      #target: loggers.WandbLogger
       params:
           name: nowname
           save_dir: "/data2/tmp/diff_log/"

examples/images/diffusion/ldm/models/autoencoder.py

@@ -6,11 +6,10 @@
 
 import torch.nn.functional as F
 from contextlib import contextmanager
+from torch.nn import Identity
 
 from ldm.modules.diffusionmodules.model import Encoder, Decoder
 from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
-
-from ldm.util import instantiate_from_config
 from ldm.modules.ema import LitEma
 
 
@@ -32,7 +31,7 @@ def __init__(self,
         self.image_key = image_key
         self.encoder = Encoder(**ddconfig)
         self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
+        self.loss = Identity(**lossconfig.get("params", dict()))
         assert ddconfig["double_z"]
         self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
         self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)

examples/images/diffusion/ldm/models/diffusion/classifier.py

@@ -9,9 +9,10 @@
 from einops import rearrange
 from glob import glob
 from natsort import natsorted
-
+from ldm.models.diffusion.ddpm import LatentDiffusion
+from ldm.lr_scheduler import LambdaLinearScheduler
 from ldm.modules.diffusionmodules.openaimodel import EncoderUNetModel, UNetModel
-from ldm.util import log_txt_as_img, default, ismap, instantiate_from_config
+from ldm.util import log_txt_as_img, default, ismap
 
 __models__ = {
     'class_label': EncoderUNetModel,
@@ -86,7 +87,7 @@ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
             print(f"Unexpected Keys: {unexpected}")
 
     def load_diffusion(self):
-        model = instantiate_from_config(self.diffusion_config)
+        model = LatentDiffusion(**self.diffusion_config.get('params',dict()))
         self.diffusion_model = model.eval()
         self.diffusion_model.train = disabled_train
         for param in self.diffusion_model.parameters():
@@ -221,7 +222,7 @@ def configure_optimizers(self):
         optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
 
         if self.use_scheduler:
-            scheduler = instantiate_from_config(self.scheduler_config)
+            scheduler = LambdaLinearScheduler(**self.scheduler_config.get('params',dict()))
 
             print("Setting up LambdaLR scheduler...")
             scheduler = [