add pl

app.py

@@ -14,11 +14,11 @@ import torch
 from omegaconf import OmegaConf
 from PIL import Image
 import spaces
-print("pip import done")
 
 from cldm.model import create_model
 from cldm.plms_hacked import PLMSSampler
 from utils_stableviton import get_batch, get_mask_location, tensor2img
+print("pip import done")
 
 PROJECT_ROOT = Path(__file__).absolute().parents[1].absolute()
 sys.path.insert(0, str(PROJECT_ROOT))

ldm/models/autoencoder.py

@@ -1,5 +1,5 @@
 import torch
-# import pytorch_lightning as pl
+import pytorch_lightning as pl
 import torch.nn as nn
 import torch.nn.functional as F
 from contextlib import contextmanager
@@ -11,7 +11,7 @@ from ldm.util import instantiate_from_config
 from ldm.modules.ema import LitEma
  
 
-class AutoencoderKL(nn.Module):
+class AutoencoderKL(pl.LightningModule):
     def __init__(self,
                  ddconfig,
                  lossconfig,

ldm/models/diffusion/ddpm.py

@@ -9,7 +9,7 @@ https://github.com/CompVis/taming-transformers
 import torch
 import torch.nn as nn
 import numpy as np
-# import pytorch_lightning as pl
+import pytorch_lightning as pl
 from torch.optim.lr_scheduler import LambdaLR
 from einops import rearrange, repeat
 from contextlib import contextmanager, nullcontext
@@ -47,7 +47,7 @@ def disabled_train(self, mode=True):
 def uniform_on_device(r1, r2, shape, device):
     return (r1 - r2) * torch.rand(*shape, device=device) + r2
 
-class DDPM(nn.Module):
+class DDPM(pl.LightningModule):
     # classic DDPM with Gaussian diffusion, in image space
     def __init__(self,
                  unet_config,

requirements.txt

@@ -20,4 +20,5 @@ cloudpickle
 fvcore
 omegaconf==2.1
 hydra-core
-pycocotools
+pycocotools
+pytorch-lightning==1.5.0