From 79342458359bf07d6241766313005386d39a3c1a Mon Sep 17 00:00:00 2001
From: Tim Dockhorn <37085732+timudk@users.noreply.github.com>
Date: Wed, 26 Jul 2023 04:26:28 -0400
Subject: [PATCH] Revert "Minimize re-exports from `__init__` files (#44)"
 (#63)

This reverts commit 57862fb4c76e72b0b867ebda3a435f0cefd71b4f.
---
 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, 14 insertions(+), 7 deletions(-)

diff --git a/sgm/__init__.py b/sgm/__init__.py
index f102a9c..f639416 100644
--- a/sgm/__init__.py
+++ b/sgm/__init__.py
@@ -1 +1,5 @@
+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 e69de29..7664a25 100644
--- a/sgm/data/__init__.py
+++ b/sgm/data/__init__.py
@@ -0,0 +1 @@
+from .dataset import StableDataModuleFromConfig
diff --git a/sgm/models/__init__.py b/sgm/models/__init__.py
index e69de29..c410b37 100644
--- a/sgm/models/__init__.py
+++ b/sgm/models/__init__.py
@@ -0,0 +1,2 @@
+from .autoencoder import AutoencodingEngine
+from .diffusion import DiffusionEngine
diff --git a/sgm/modules/__init__.py b/sgm/modules/__init__.py
index ba75205..0db1d77 100644
--- a/sgm/modules/__init__.py
+++ b/sgm/modules/__init__.py
@@ -1,12 +1,5 @@
 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 e69de29..ce7968a 100644
--- a/sgm/modules/diffusionmodules/__init__.py
+++ b/sgm/modules/diffusionmodules/__init__.py
@@ -0,0 +1,7 @@
+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