Update app.py

app.py

@@ -33,7 +33,6 @@ import torchvision
 from torchvision import transforms
 from kornia.filters import filter2d
 
-from lightweight_gan.diff_augment import DiffAugment
 from lightweight_gan.version import __version__
 
 from tqdm import tqdm
@@ -43,6 +42,13 @@ from adabelief_pytorch import AdaBelief
 
 # helpers
 
+
+def DiffAugment(x, types=[]):
+    for p in types:
+        for f in AUGMENT_FNS[p]:
+            x = f(x)
+    return x.contiguous()
+
 def exists(val):
     return val is not None