Update imports (#279)

marksgraham · web-flow · commit 4849c6b40940 · 2023-02-27T21:56:33.000Z
diff --git a/generative/losses/__init__.py b/generative/losses/__init__.py
@@ -11,4 +11,6 @@
 
 from __future__ import annotations
 
+from .adversarial_loss import PatchAdversarialLoss
+from .perceptual import PerceptualLoss
 from .spectral_loss import JukeboxLoss
diff --git a/tests/test_adversarial.py b/tests/test_adversarial.py
@@ -16,7 +16,7 @@
 import torch
 from parameterized import parameterized
 
-from generative.losses.adversarial_loss import PatchAdversarialLoss
+from generative.losses import PatchAdversarialLoss
 
 shapes_tensors = {"2d": [4, 1, 64, 64], "3d": [4, 1, 64, 64, 64]}
 reductions = ["sum", "mean"]
diff --git a/tests/test_perceptual_loss.py b/tests/test_perceptual_loss.py
@@ -16,7 +16,7 @@
 import torch
 from parameterized import parameterized
 
-from generative.losses.perceptual import PerceptualLoss
+from generative.losses import PerceptualLoss
 
 TEST_CASES = [
     [{"spatial_dims": 2, "network_type": "squeeze"}, (2, 1, 64, 64), (2, 1, 64, 64)],
diff --git a/tutorials/generative/2d_autoencoderkl/2d_autoencoderkl_tutorial.ipynb b/tutorials/generative/2d_autoencoderkl/2d_autoencoderkl_tutorial.ipynb
@@ -120,8 +120,7 @@
     "from torch.nn import L1Loss\n",
     "from tqdm import tqdm\n",
     "\n",
-    "from generative.losses.adversarial_loss import PatchAdversarialLoss\n",
-    "from generative.losses.perceptual import PerceptualLoss\n",
+    "from generative.losses import PatchAdversarialLoss, PerceptualLoss\n",
     "from generative.networks.nets import AutoencoderKL, PatchDiscriminator\n",
     "\n",
     "print_config()"
diff --git a/tutorials/generative/2d_autoencoderkl/2d_autoencoderkl_tutorial.py b/tutorials/generative/2d_autoencoderkl/2d_autoencoderkl_tutorial.py
@@ -47,8 +47,7 @@
 from torch.nn import L1Loss
 from tqdm import tqdm
 
-from generative.losses.adversarial_loss import PatchAdversarialLoss
-from generative.losses.perceptual import PerceptualLoss
+from generative.losses import PatchAdversarialLoss, PerceptualLoss
 from generative.networks.nets import AutoencoderKL, PatchDiscriminator
 
 print_config()
diff --git a/tutorials/generative/2d_super_resolution/2d_stable_diffusion_v2_super_resolution.ipynb b/tutorials/generative/2d_super_resolution/2d_stable_diffusion_v2_super_resolution.ipynb
@@ -124,8 +124,7 @@
     "from torch.cuda.amp import GradScaler, autocast\n",
     "from tqdm import tqdm\n",
     "\n",
-    "from generative.losses.adversarial_loss import PatchAdversarialLoss\n",
-    "from generative.losses.perceptual import PerceptualLoss\n",
+    "from generative.losses import PatchAdversarialLoss, PerceptualLoss\n",
     "from generative.networks.nets import AutoencoderKL, DiffusionModelUNet, PatchDiscriminator\n",
     "from generative.networks.schedulers import DDPMScheduler\n",
     "\n",
diff --git a/tutorials/generative/2d_super_resolution/2d_stable_diffusion_v2_super_resolution.py b/tutorials/generative/2d_super_resolution/2d_stable_diffusion_v2_super_resolution.py
@@ -62,8 +62,7 @@
 from torch.cuda.amp import GradScaler, autocast
 from tqdm import tqdm
 
-from generative.losses.adversarial_loss import PatchAdversarialLoss
-from generative.losses.perceptual import PerceptualLoss
+from generative.losses import PatchAdversarialLoss, PerceptualLoss
 from generative.networks.nets import AutoencoderKL, DiffusionModelUNet, PatchDiscriminator
 from generative.networks.schedulers import DDPMScheduler
 
diff --git a/tutorials/generative/2d_vqgan/2d_vqgan_tutorial.ipynb b/tutorials/generative/2d_vqgan/2d_vqgan_tutorial.ipynb
@@ -103,8 +103,7 @@
     "from tqdm import tqdm\n",
     "\n",
     "# TODO: Add right import reference after deployed\n",
-    "from generative.losses.adversarial_loss import PatchAdversarialLoss\n",
-    "from generative.losses.perceptual import PerceptualLoss\n",
+    "from generative.losses import PatchAdversarialLoss, PerceptualLoss\n",
     "from generative.networks.nets import VQVAE, PatchDiscriminator\n",
     "\n",
     "print_config()"
diff --git a/tutorials/generative/2d_vqgan/2d_vqgan_tutorial.py b/tutorials/generative/2d_vqgan/2d_vqgan_tutorial.py
@@ -59,8 +59,7 @@
 from tqdm import tqdm
 
 # TODO: Add right import reference after deployed
-from generative.losses.adversarial_loss import PatchAdversarialLoss
-from generative.losses.perceptual import PerceptualLoss
+from generative.losses import PatchAdversarialLoss, PerceptualLoss
 from generative.networks.nets import VQVAE, PatchDiscriminator
 
 print_config()
diff --git a/tutorials/generative/3d_autoencoderkl/3d_autoencoderkl_tutorial.ipynb b/tutorials/generative/3d_autoencoderkl/3d_autoencoderkl_tutorial.ipynb
@@ -118,8 +118,7 @@
     "from torch.cuda.amp import autocast\n",
     "from tqdm import tqdm\n",
     "\n",
-    "from generative.losses.adversarial_loss import PatchAdversarialLoss\n",
-    "from generative.losses.perceptual import PerceptualLoss\n",
+    "from generative.losses import PatchAdversarialLoss, PerceptualLoss\n",
     "from generative.networks.nets import AutoencoderKL, PatchDiscriminator\n",
     "\n",
     "print_config()"
diff --git a/tutorials/generative/3d_autoencoderkl/3d_autoencoderkl_tutorial.py b/tutorials/generative/3d_autoencoderkl/3d_autoencoderkl_tutorial.py
@@ -51,8 +51,7 @@
 from torch.cuda.amp import autocast
 from tqdm import tqdm
 
-from generative.losses.adversarial_loss import PatchAdversarialLoss
-from generative.losses.perceptual import PerceptualLoss
+from generative.losses import PatchAdversarialLoss, PerceptualLoss
 from generative.networks.nets import AutoencoderKL, PatchDiscriminator
 
 print_config()
diff --git a/tutorials/generative/3d_ldm/3d_ldm_tutorial.ipynb b/tutorials/generative/3d_ldm/3d_ldm_tutorial.ipynb
@@ -83,8 +83,7 @@
     "from tqdm import tqdm\n",
     "\n",
     "from generative.inferers import LatentDiffusionInferer\n",
-    "from generative.losses.adversarial_loss import PatchAdversarialLoss\n",
-    "from generative.losses.perceptual import PerceptualLoss\n",
+    "from generative.losses import PatchAdversarialLoss, PerceptualLoss\n",
     "from generative.networks.nets import AutoencoderKL, DiffusionModelUNet, PatchDiscriminator\n",
     "from generative.networks.schedulers import DDPMScheduler\n",
     "\n",
diff --git a/tutorials/generative/3d_ldm/3d_ldm_tutorial.py b/tutorials/generative/3d_ldm/3d_ldm_tutorial.py
@@ -36,8 +36,7 @@
 from tqdm import tqdm
 
 from generative.inferers import LatentDiffusionInferer
-from generative.losses.adversarial_loss import PatchAdversarialLoss
-from generative.losses.perceptual import PerceptualLoss
+from generative.losses import PatchAdversarialLoss, PerceptualLoss
 from generative.networks.nets import AutoencoderKL, DiffusionModelUNet, PatchDiscriminator
 from generative.networks.schedulers import DDPMScheduler
 
diff --git a/tutorials/generative/classifier_free_guidance/2d_ddpm_classifier_free_guidance_tutorial.ipynb b/tutorials/generative/classifier_free_guidance/2d_ddpm_classifier_free_guidance_tutorial.ipynb
@@ -115,7 +115,7 @@
     "\n",
     "# TODO: Add right import reference after deployed\n",
     "from generative.networks.nets import DiffusionModelUNet\n",
-    "from generative.schedulers import DDPMScheduler\n",
+    "from generative.networks.schedulers import DDPMScheduler\n",
     "\n",
     "print_config()"
    ]
diff --git a/tutorials/generative/classifier_free_guidance/2d_ddpm_classifier_free_guidance_tutorial.py b/tutorials/generative/classifier_free_guidance/2d_ddpm_classifier_free_guidance_tutorial.py
@@ -67,7 +67,7 @@
 
 # TODO: Add right import reference after deployed
 from generative.networks.nets import DiffusionModelUNet
-from generative.schedulers import DDPMScheduler
+from generative.networks.schedulers import DDPMScheduler
 
 print_config()
 
