diff --git a/setup.py b/setup.py
index dbe8ad7..9b57672 100644
--- a/setup.py
+++ b/setup.py
@@ -74,7 +74,7 @@ setup(
     description="Steganography tool based on DeepLearning GANs",
     entry_points={
         'console_scripts': [
-            'steganogan=steganogan:cli.main'
+            'steganogan=steganogan.cli:main'
         ],
     },
     extras_require={
diff --git a/steganogan/__init__.py b/steganogan/__init__.py
index be09288..afe89d7 100644
--- a/steganogan/__init__.py
+++ b/steganogan/__init__.py
@@ -5,7 +5,6 @@ __author__ = """MIT Data To AI Lab"""
 __email__ = 'dailabmit@gmail.com'
 __version__ = '0.1.2-dev'
 
-from steganogan import cli
 from steganogan.models import SteganoGAN
 
-__all__ = ('SteganoGAN', 'cli')
+__all__ = ('SteganoGAN', )
diff --git a/steganogan/cli.py b/steganogan/cli.py
index 1ae4d3c..6a833fd 100644
--- a/steganogan/cli.py
+++ b/steganogan/cli.py
@@ -1,22 +1,23 @@
 # -*- coding: utf-8 -*-
 
-"""Top-level package for SteganoGAN."""
-
-__author__ = 'MIT Data To AI Lab'
-__email__ = 'dailabmit@gmail.com'
-__version__ = '0.1.0.dev.dev'
-
 import argparse
-import os
 
 from steganogan.models import SteganoGAN
 
 
 def _get_steganogan(args):
-    model_name = '{}.steg'.format(args.architecture)
-    pretrained_path = os.path.join(os.path.dirname(__file__), 'pretrained')
-    model_path = os.path.join(pretrained_path, model_name)
-    return SteganoGAN.load(model_path, cuda=not args.cpu, verbose=args.verbose)
+
+    steganogan_kwargs = {
+        'cuda': not args.cpu,
+        'verbose': args.verbose
+    }
+
+    if args.path:
+        steganogan_kwargs['path'] = args.path
+    else:
+        steganogan_kwargs['architecture'] = args.architecture
+
+    return SteganoGAN.load(**steganogan_kwargs)
 
 
 def _encode(args):
@@ -44,8 +45,12 @@ def _get_parser():
     # Parent Parser - Shared options
     parent = argparse.ArgumentParser(add_help=False)
     parent.add_argument('-v', '--verbose', action='store_true', help='Be verbose')
-    parent.add_argument('-a', '--architecture', default='dense', choices=('basic', 'dense'),
-                        help='Model architecture. Use the same one for both encoding and decoding')
+    group = parent.add_mutually_exclusive_group()
+    group.add_argument('-a', '--architecture', default='dense',
+                       choices={'basic', 'dense', 'residual'},
+                       help='Model architecture. Use the same one for both encoding and decoding')
+
+    group.add_argument('-p', '--path', help='Load a pretrained model from a given path.')
     parent.add_argument('--cpu', action='store_true',
                         help='Force CPU usage even if CUDA is available')
 
diff --git a/steganogan/models.py b/steganogan/models.py
index 78afdbc..f236e3e 100644
--- a/steganogan/models.py
+++ b/steganogan/models.py
@@ -341,8 +341,26 @@ class SteganoGAN(object):
         torch.save(self, path)
 
     @classmethod
-    def load(cls, path, cuda=True, verbose=False):
-        """Loads an instance of SteganoGAN from the given path."""
+    def load(cls, architecture=None, path=None, cuda=True, verbose=False):
+        """Loads an instance of SteganoGAN for the given architecture (default pretrained models)
+        or loads a pretrained model from a given path.
+
+        Args:
+            architecture(str): Name of a pretrained model to be loaded from the default models.
+            path(str): Path to custom pretrained model. *Architecture must be None.
+            cuda(bool): Force loaded model to use cuda (if available).
+            verbose(bool): Force loaded model to use or not verbose.
+        """
+
+        if architecture and not path:
+            model_name = '{}.steg'.format(architecture)
+            pretrained_path = os.path.join(os.path.dirname(__file__), 'pretrained')
+            path = os.path.join(pretrained_path, model_name)
+
+        elif (architecture is None and path is None) or (architecture and path):
+            raise ValueError(
+                'Please provide either an architecture or a path to pretrained model.')
+
         steganogan = torch.load(path, map_location='cpu')
         steganogan.verbose = verbose
 

Move `_get_steganogan` logic to `SteganoGAN.load`
The code that loads the pretrained models is currently inside a "private" function in the `cli.py` module.

This should be moved inside the `SteganoGAN.load` method to allow loading pretrained models directly from the class.

For this, the following changes are needed:

* Add an `architecture` argument to the `load` method.
* Make the `path` argument from the `load` method optional (`=None`).

Then, reimplement the load method so that:

* If a path is given, it is directly used to load the model from it.
* If an architecture name is given, the path to the model is built dynamically in a similar way to what we currently have in `_get_stegangan`.
* If either both arguments or none are given, an exception is raised.

Then, in the `cli.py` module, the _get_steganogan implementation can be changed pass the architecture value directly to the `load` method. 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index db415fe..b33945e 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,13 +1,75 @@
 # -*- coding: utf-8 -*-
 
-import os
 from unittest.mock import MagicMock, patch
 
 from steganogan import cli
 
 
 @patch('steganogan.cli.SteganoGAN.load')
-def test__get_steganogan(mock_steganogan_load):
+def test__get_steganogan_with_path_and_architecture(mock_steganogan_load):
+    """
+    Test that:
+        * The model is called with the path and not acrhitecture.
+        * SteganoGAN.load is called with the right values
+        * The output of SteganoGAN.load is returned
+    """
+
+    # setup
+    mock_steganogan_load.return_value = 'Steganogan'
+
+    params = MagicMock(
+        architecture='dense',
+        cpu=True,
+        path='my_path/basic',
+        verbose=True,
+    )
+
+    # run
+    cli_test = cli._get_steganogan(params)
+
+    # assert
+    mock_steganogan_load.assert_called_once_with(
+        path='my_path/basic',
+        cuda=False,
+        verbose=True
+    )
+
+    assert cli_test == 'Steganogan'
+
+
+@patch('steganogan.cli.SteganoGAN.load')
+def test__get_steganogan_with_path(mock_steganogan_load):
+    """
+    Test that:
+        * The model is loaded with the path.
+        * SteganoGAN.load is called with the right values
+        * The output of SteganoGAN.load is returned
+    """
+
+    # setup
+    mock_steganogan_load.return_value = 'Steganogan'
+
+    params = MagicMock(
+        cpu=True,
+        path='my_path/basic',
+        verbose=True,
+    )
+
+    # run
+    cli_test = cli._get_steganogan(params)
+
+    # assert
+    mock_steganogan_load.assert_called_once_with(
+        path='my_path/basic',
+        cuda=False,
+        verbose=True
+    )
+
+    assert cli_test == 'Steganogan'
+
+
+@patch('steganogan.cli.SteganoGAN.load')
+def test__get_steganogan_with_architecture(mock_steganogan_load):
     """
     Test that:
         * The model path is the right one, with the right architecture
@@ -22,20 +84,15 @@ def test__get_steganogan(mock_steganogan_load):
         cpu=True,
         architecture='basic',
         verbose=True,
+        path=None
     )
 
-    model_name = '{}.steg'.format(params.architecture)
-    parent_path = os.path.dirname(os.path.dirname(__file__))
-    stega_path = os.path.join(parent_path, 'steganogan')
-    pretrained_path = os.path.join(stega_path, 'pretrained')
-    model_path = os.path.join(pretrained_path, model_name)
-
     # run
     cli_test = cli._get_steganogan(params)
 
     # assert
     mock_steganogan_load.assert_called_once_with(
-        model_path,
+        architecture='basic',
         cuda=False,
         verbose=True
     )
diff --git a/tests/test_models.py b/tests/test_models.py
index c3b9877..e4185f6 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1067,19 +1067,20 @@ class TestSteganoGAN(TestCase):
         # assert
         mock_save.assert_called_once_with(steganogan, 'some_path')
 
+    @patch('steganogan.models.os.path.join')
     @patch('steganogan.models.torch.load')
-    def test_load(self, mock_load):
-        """Test when loading a path, we execute and update steganogan"""
+    def test_load(self, mock_load, os_path_mock):
+        """Test loading a default architecture"""
 
         # setup
         steganogan = MagicMock()
         mock_load.return_value = steganogan
 
         # run
-        result = models.SteganoGAN.load('some_path', cuda=False, verbose=False)
+        result = models.SteganoGAN.load('some_architecture', cuda=False, verbose=False)
 
         # assert
-        mock_load.assert_called_once_with('some_path', map_location='cpu')
+        mock_load.assert_called_once_with(os_path_mock.return_value, map_location='cpu')
 
         assert not steganogan.verbose
 
@@ -1091,8 +1092,9 @@ class TestSteganoGAN(TestCase):
 
         assert result == steganogan
 
+    @patch('steganogan.models.os.path.join')
     @patch('steganogan.models.torch.load')
-    def test_load_cuda_verbose_true(self, mock_load):
+    def test_load_cuda_verbose_true(self, mock_load, os_path_mock):
         """Test loading a path, with cuda and verbose True"""
 
         # setup
@@ -1100,10 +1102,10 @@ class TestSteganoGAN(TestCase):
         mock_load.return_value = steganogan
 
         # run
-        result = models.SteganoGAN.load('some_path', cuda=True, verbose=True)
+        result = models.SteganoGAN.load(architecture='some_architecture', cuda=True, verbose=True)
 
         # assert
-        mock_load.assert_called_once_with('some_path', map_location='cpu')
+        mock_load.assert_called_once_with(os_path_mock.return_value, map_location='cpu')
 
         assert steganogan.verbose
 
@@ -1114,3 +1116,26 @@ class TestSteganoGAN(TestCase):
         steganogan.set_device.assert_called_once_with(True)
 
         assert result == steganogan
+
+    @patch('steganogan.models.torch.load')
+    def test_load_path_no_architecture(self, mock_load):
+        """Test loading a model when passing a path and architecture is None"""
+
+        # setup
+        steganogan = MagicMock()
+        mock_load.return_value = steganogan
+
+        # run
+        result = models.SteganoGAN.load(
+            architecture=None, path='some_path', cuda=True, verbose=True)
+
+        # assert
+        mock_load.assert_called_once_with('some_path', map_location='cpu')
+
+        assert result == steganogan
+
+    def test_load_path_and_architecture(self):
+
+        # run / assert
+        with self.assertRaises(ValueError):
+            models.SteganoGAN.load(architecture='some_arch', path='some_path')

pip install -e .[dev]
parse_log_pytest
requirements.txt
pytest
3.6
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
alabaster==0.7.13
argon2-cffi==21.3.0
argon2-cffi-bindings==21.2.0
async-generator==1.10
attrs==22.2.0
autoflake==1.4
autopep8==2.0.0
Babel==2.11.0
backcall==0.2.0
bleach==4.1.0
bump2version==1.0.1
bumpversion==0.6.0
certifi==2021.5.30
cffi==1.15.1
charset-normalizer==2.0.12
colorama==0.4.5
comm==0.1.4
coverage==6.2
cryptography==40.0.2
dataclasses==0.8
decorator==5.1.1
defusedxml==0.7.1
distlib==0.3.9
docutils==0.17.1
entrypoints==0.4
filelock==3.4.1
flake8==5.0.4
idna==3.10
imageio==2.15.0
imagesize==1.4.1
importlib-metadata==4.2.0
importlib-resources==5.4.0
iniconfig==1.1.1
ipykernel==5.5.6
ipython==7.16.3
ipython-genutils==0.2.0
ipywidgets==7.8.5
isort==5.10.1
jedi==0.17.2
jeepney==0.7.1
Jinja2==3.0.3
jsonschema==3.2.0
jupyter==1.0.0
jupyter-client==7.1.2
jupyter-console==6.4.3
jupyter-core==4.9.2
jupyterlab-pygments==0.1.2
jupyterlab_widgets==1.1.11
keyring==23.4.1
m2r==0.3.1
MarkupSafe==2.0.1
mccabe==0.7.0
mistune==0.8.4
nbclient==0.5.9
nbconvert==6.0.7
nbformat==5.1.3
nest-asyncio==1.6.0
notebook==6.4.10
numpy==1.19.5
packaging==21.3
pandocfilters==1.5.1
parso==0.7.1
pexpect==4.9.0
pickleshare==0.7.5
Pillow==8.4.0
pkginfo==1.10.0
platformdirs==2.4.0
pluggy==1.0.0
prometheus-client==0.17.1
prompt-toolkit==3.0.36
ptyprocess==0.7.0
py==1.11.0
pycodestyle==2.9.1
pycparser==2.21
pyflakes==2.5.0
Pygments==2.14.0
pyparsing==3.1.4
pyrsistent==0.18.0
pytest==7.0.1
python-dateutil==2.9.0.post0
pytz==2025.2
pyzmq==25.1.2
qtconsole==5.2.2
QtPy==2.0.1
readme-renderer==34.0
reedsolo==1.7.0
requests==2.27.1
requests-toolbelt==1.0.0
rfc3986==1.5.0
scipy==1.5.4
SecretStorage==3.3.3
Send2Trash==1.8.3
six==1.17.0
snowballstemmer==2.2.0
Sphinx==4.3.2
sphinx-rtd-theme==1.3.0
sphinxcontrib-applehelp==1.0.2
sphinxcontrib-devhelp==1.0.2
sphinxcontrib-htmlhelp==2.0.0
sphinxcontrib-jquery==4.1
sphinxcontrib-jsmath==1.0.1
sphinxcontrib-qthelp==1.0.3
sphinxcontrib-serializinghtml==1.1.5
-e git+https://github.com/DAI-Lab/SteganoGAN.git@4b0e6019078e99e36a8b095636620b107f3fa3d1#egg=steganogan
terminado==0.12.1
testpath==0.6.0
toml==0.10.2
tomli==1.2.3
torch==1.10.1
torchvision==0.11.2
tornado==6.1
tox==3.28.0
tqdm==4.64.1
traitlets==4.3.3
twine==3.8.0
typing_extensions==4.1.1
urllib3==1.26.20
virtualenv==20.16.2
watchdog==2.3.1
wcwidth==0.2.13
webencodings==0.5.1
widgetsnbextension==3.6.10
zipp==3.6.0

name: SteganoGAN
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - certifi=2021.5.30=py36h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.3=he6710b0_2
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=1.1.1w=h7f8727e_0
  - pip=21.2.2=py36h06a4308_0
  - python=3.6.13=h12debd9_1
  - readline=8.2=h5eee18b_0
  - setuptools=58.0.4=py36h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - wheel=0.37.1=pyhd3eb1b0_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - alabaster==0.7.13
      - argon2-cffi==21.3.0
      - argon2-cffi-bindings==21.2.0
      - async-generator==1.10
      - attrs==22.2.0
      - autoflake==1.4
      - autopep8==2.0.0
      - babel==2.11.0
      - backcall==0.2.0
      - bleach==4.1.0
      - bump2version==1.0.1
      - bumpversion==0.6.0
      - cffi==1.15.1
      - charset-normalizer==2.0.12
      - colorama==0.4.5
      - comm==0.1.4
      - coverage==6.2
      - cryptography==40.0.2
      - dataclasses==0.8
      - decorator==5.1.1
      - defusedxml==0.7.1
      - distlib==0.3.9
      - docutils==0.17.1
      - entrypoints==0.4
      - filelock==3.4.1
      - flake8==5.0.4
      - idna==3.10
      - imageio==2.15.0
      - imagesize==1.4.1
      - importlib-metadata==4.2.0
      - importlib-resources==5.4.0
      - iniconfig==1.1.1
      - ipykernel==5.5.6
      - ipython==7.16.3
      - ipython-genutils==0.2.0
      - ipywidgets==7.8.5
      - isort==5.10.1
      - jedi==0.17.2
      - jeepney==0.7.1
      - jinja2==3.0.3
      - jsonschema==3.2.0
      - jupyter==1.0.0
      - jupyter-client==7.1.2
      - jupyter-console==6.4.3
      - jupyter-core==4.9.2
      - jupyterlab-pygments==0.1.2
      - jupyterlab-widgets==1.1.11
      - keyring==23.4.1
      - m2r==0.3.1
      - markupsafe==2.0.1
      - mccabe==0.7.0
      - mistune==0.8.4
      - nbclient==0.5.9
      - nbconvert==6.0.7
      - nbformat==5.1.3
      - nest-asyncio==1.6.0
      - notebook==6.4.10
      - numpy==1.19.5
      - packaging==21.3
      - pandocfilters==1.5.1
      - parso==0.7.1
      - pexpect==4.9.0
      - pickleshare==0.7.5
      - pillow==8.4.0
      - pkginfo==1.10.0
      - platformdirs==2.4.0
      - pluggy==1.0.0
      - prometheus-client==0.17.1
      - prompt-toolkit==3.0.36
      - ptyprocess==0.7.0
      - py==1.11.0
      - pycodestyle==2.9.1
      - pycparser==2.21
      - pyflakes==2.5.0
      - pygments==2.14.0
      - pyparsing==3.1.4
      - pyrsistent==0.18.0
      - pytest==7.0.1
      - python-dateutil==2.9.0.post0
      - pytz==2025.2
      - pyzmq==25.1.2
      - qtconsole==5.2.2
      - qtpy==2.0.1
      - readme-renderer==34.0
      - reedsolo==1.7.0
      - requests==2.27.1
      - requests-toolbelt==1.0.0
      - rfc3986==1.5.0
      - scipy==1.5.4
      - secretstorage==3.3.3
      - send2trash==1.8.3
      - six==1.17.0
      - snowballstemmer==2.2.0
      - sphinx==4.3.2
      - sphinx-rtd-theme==1.3.0
      - sphinxcontrib-applehelp==1.0.2
      - sphinxcontrib-devhelp==1.0.2
      - sphinxcontrib-htmlhelp==2.0.0
      - sphinxcontrib-jquery==4.1
      - sphinxcontrib-jsmath==1.0.1
      - sphinxcontrib-qthelp==1.0.3
      - sphinxcontrib-serializinghtml==1.1.5
      - terminado==0.12.1
      - testpath==0.6.0
      - toml==0.10.2
      - tomli==1.2.3
      - torch==1.10.1
      - torchvision==0.11.2
      - tornado==6.1
      - tox==3.28.0
      - tqdm==4.64.1
      - traitlets==4.3.3
      - twine==3.8.0
      - typing-extensions==4.1.1
      - urllib3==1.26.20
      - virtualenv==20.16.2
      - watchdog==2.3.1
      - wcwidth==0.2.13
      - webencodings==0.5.1
      - widgetsnbextension==3.6.10
      - zipp==3.6.0
prefix: /opt/conda/envs/SteganoGAN

tests/test_cli.py::test__get_steganogan_with_path_and_architecture tests/test_cli.py::test__get_steganogan_with_path tests/test_cli.py::test__get_steganogan_with_architecture tests/test_models.py::TestSteganoGAN::test_load tests/test_models.py::TestSteganoGAN::test_load_cuda_verbose_true tests/test_models.py::TestSteganoGAN::test_load_path_and_architecture tests/test_models.py::TestSteganoGAN::test_load_path_no_architecture
tests/test_models.py::TestSteganoGAN::test__critic
tests/test_cli.py::test__encode tests/test_cli.py::test__decode tests/test_models.py::TestSteganoGAN::test___init__log_dir tests/test_models.py::TestSteganoGAN::test___init__without_logdir tests/test_models.py::TestSteganoGAN::test__coding_scores tests/test_models.py::TestSteganoGAN::test__encode_decode_quantize_false tests/test_models.py::TestSteganoGAN::test__encode_decode_quantize_true tests/test_models.py::TestSteganoGAN::test__fit_coders tests/test_models.py::TestSteganoGAN::test__fit_critic tests/test_models.py::TestSteganoGAN::test__generate_samples tests/test_models.py::TestSteganoGAN::test__get_instance_is_class tests/test_models.py::TestSteganoGAN::test__get_instance_is_instance tests/test_models.py::TestSteganoGAN::test__get_optimizers tests/test_models.py::TestSteganoGAN::test__make_payload tests/test_models.py::TestSteganoGAN::test__random_data tests/test_models.py::TestSteganoGAN::test__validate tests/test_models.py::TestSteganoGAN::test_decode_image tests/test_models.py::TestSteganoGAN::test_decode_image_is_not_file tests/test_models.py::TestSteganoGAN::test_decode_image_zero_candidates tests/test_models.py::TestSteganoGAN::test_encode tests/test_models.py::TestSteganoGAN::test_fit_optimizer_is_none tests/test_models.py::TestSteganoGAN::test_fit_with_cuda tests/test_models.py::TestSteganoGAN::test_fit_with_log_dir tests/test_models.py::TestSteganoGAN::test_fit_with_optimizers tests/test_models.py::TestSteganoGAN::test_save tests/test_models.py::TestSteganoGAN::test_set_device_cpu tests/test_models.py::TestSteganoGAN::test_set_device_cuda
setup.py steganogan/models.py steganogan/cli.py steganogan/__init__.py
setup.py steganogan/models.py steganogan/cli.py steganogan/__init__.py