From 807ca115da17ffd682c6ec5f74ef41afd2d87483 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Mon, 17 Jul 2023 10:04:14 +0300
Subject: [PATCH] Minimize re-exports from `__init__` files

This allows importing parts of the package without having to
import practically everything (since importing a package will
import its parents' __init__s, etc).
---
 sgm/__init__.py                          | 4 ----
 sgm/data/__init__.py                     | 1 -
 sgm/models/__init__.py                   | 2 --
 sgm/modules/__init__.py                  | 7 +++++++
 sgm/modules/diffusionmodules/__init__.py | 7 -------
 5 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/sgm/__init__.py b/sgm/__init__.py
index f639416e..f102a9ca 100644
--- a/sgm/__init__.py
+++ b/sgm/__init__.py
@@ -1,5 +1 @@
-from .data import StableDataModuleFromConfig
-from .models import AutoencodingEngine, DiffusionEngine
-from .util import instantiate_from_config, get_configs_path
-
 __version__ = "0.0.1"
diff --git a/sgm/data/__init__.py b/sgm/data/__init__.py
index 7664a25c..e69de29b 100644
--- a/sgm/data/__init__.py
+++ b/sgm/data/__init__.py
@@ -1 +0,0 @@
-from .dataset import StableDataModuleFromConfig
diff --git a/sgm/models/__init__.py b/sgm/models/__init__.py
index c410b374..e69de29b 100644
--- a/sgm/models/__init__.py
+++ b/sgm/models/__init__.py
@@ -1,2 +0,0 @@
-from .autoencoder import AutoencodingEngine
-from .diffusion import DiffusionEngine
diff --git a/sgm/modules/__init__.py b/sgm/modules/__init__.py
index 0db1d771..ba75205c 100644
--- a/sgm/modules/__init__.py
+++ b/sgm/modules/__init__.py
@@ -1,5 +1,12 @@
 from .encoders.modules import GeneralConditioner
 
+__all__ = [
+    # `sgm.models.GeneralConditioner` is referenced in model configurations, etc.,
+    # so it must be re-exported from this module.
+    "GeneralConditioner",
+    "UNCONDITIONAL_CONFIG",
+]
+
 UNCONDITIONAL_CONFIG = {
     "target": "sgm.modules.GeneralConditioner",
     "params": {"emb_models": []},
diff --git a/sgm/modules/diffusionmodules/__init__.py b/sgm/modules/diffusionmodules/__init__.py
index ce7968af..e69de29b 100644
--- a/sgm/modules/diffusionmodules/__init__.py
+++ b/sgm/modules/diffusionmodules/__init__.py
@@ -1,7 +0,0 @@
-from .denoiser import Denoiser
-from .discretizer import Discretization
-from .loss import StandardDiffusionLoss
-from .model import Model, Encoder, Decoder
-from .openaimodel import UNetModel
-from .sampling import BaseDiffusionSampler
-from .wrappers import OpenAIWrapper