From 0812d6252c95f0498c46ab75d086555fa1064cbf Mon Sep 17 00:00:00 2001
From: Anthony Wu <462072+anthonywu@users.noreply.github.com>
Date: Mon, 16 Sep 2024 10:27:04 -0700
Subject: [PATCH] simplify import paths of frequently used items in scripts and
 demos

---
 README.md                    | 3 +--
 src/mflux/__init__.py        | 5 +++++
 src/mflux/generate.py        | 4 +---
 src/mflux/save.py            | 3 +--
 tests/test_readme_example.sh | 7 +++++++
 5 files changed, 15 insertions(+), 7 deletions(-)
 create mode 100755 tests/test_readme_example.sh

diff --git a/README.md b/README.md
index 8fd15b0..3f5539f 100644
--- a/README.md
+++ b/README.md
@@ -105,8 +105,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 Or, with the correct python environment active, create and run a separate script like the following:
 
 ```python
-from mflux.flux.flux import Flux1
-from mflux.config.config import Config
+from mflux import Flux1, Config
 
 # Load the model
 flux = Flux1.from_alias(
diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py
index e69de29..3994351 100644
--- a/src/mflux/__init__.py
+++ b/src/mflux/__init__.py
@@ -0,0 +1,5 @@
+from mflux.flux.flux import Flux1
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+
+__all__ = ["Flux1", "Config", "ModelConfig"]
diff --git a/src/mflux/generate.py b/src/mflux/generate.py
index 1b8e59f..e4a6640 100644
--- a/src/mflux/generate.py
+++ b/src/mflux/generate.py
@@ -5,9 +5,7 @@ import time
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from mflux.config.model_config import ModelConfig
-from mflux.config.config import Config
-from mflux.flux.flux import Flux1
+from mflux import Flux1, Config, ModelConfig
 
 
 def main():
diff --git a/src/mflux/save.py b/src/mflux/save.py
index e34bf8f..0bfd4bd 100644
--- a/src/mflux/save.py
+++ b/src/mflux/save.py
@@ -4,8 +4,7 @@ import sys
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from mflux.flux.flux import Flux1
-from mflux.config.model_config import ModelConfig
+from mflux import Flux1, ModelConfig
 
 
 def main():
diff --git a/tests/test_readme_example.sh b/tests/test_readme_example.sh
new file mode 100755
index 0000000..3fc6827
--- /dev/null
+++ b/tests/test_readme_example.sh
@@ -0,0 +1,7 @@
+mflux-generate \
+    --prompt "Luxury food photograph" \
+    --model schnell \
+    --steps 2 \
+    --seed 2 \
+    --height 1024 \
+    --width 1024