careamics 0.0.2__tar.gz → 0.0.3__tar.gz

{careamics-0.0.2 → careamics-0.0.3}/.pre-commit-config.yaml

@@ -30,7 +30,8 @@ repos:
     hooks:
       - id: mypy
         files: "^src/"
-        exclude: "^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*"
+        exclude: "^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*|^src/careamics/config/likelihood_model.py|^src/careamics/losses/loss_factory.py|^src/careamics/losses/lvae/losses.py"
+        args: ['--config-file', 'mypy.ini']
         additional_dependencies:
           - numpy
           - types-PyYAML
@@ -41,7 +42,7 @@ repos:
     rev: v1.8.0rc2
     hooks:
       - id: numpydoc-validation
-        exclude: "^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*"      
+        exclude: "^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*|^src/careamics/losses/lvae/.*"      
 
   # # jupyter linting and formatting
   # - repo: https://github.com/nbQA-dev/nbQA

{careamics-0.0.2 → careamics-0.0.3}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.3
 Name: careamics
-Version: 0.0.2
+Version: 0.0.3
 Summary: Toolbox for running N2V and friends.
 Project-URL: homepage, https://careamics.github.io/
 Project-URL: repository, https://github.com/CAREamics/careamics
-Author-email: Melisande Croft <melisande.croft@fht.org>, Joran Deschamps <joran.deschamps@fht.org>, Igor Zubarev <igor.zubarev@fht.org>
+Author-email: CAREamics team <rse@fht.org>, Ashesh <ashesh.ashesh@fht.org>, Federico Carrara <federico.carrara@fht.org>, Melisande Croft <melisande.croft@fht.org>, Joran Deschamps <joran.deschamps@fht.org>, Vera Galinova <vera.galinova@fht.org>, Igor Zubarev <igor.zubarev@fht.org>
 License: BSD-3-Clause
 License-File: LICENSE
 Classifier: Development Status :: 3 - Alpha

careamics-0.0.3/examples/example_training_LVAE_split.ipynb

@@ -0,0 +1,439 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import socket\n",
+    "from pathlib import Path\n",
+    "from typing import Optional, Literal, Union\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch \n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from pytorch_lightning import Trainer\n",
+    "from pytorch_lightning.loggers import WandbLogger\n",
+    "from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint, EarlyStopping\n",
+    "\n",
+    "from careamics.config import VAEAlgorithmConfig\n",
+    "from careamics.config.architectures import LVAEModel\n",
+    "from careamics.config.callback_model import CheckpointModel, EarlyStoppingModel\n",
+    "from careamics.config.likelihood_model import (\n",
+    "    GaussianLikelihoodConfig,\n",
+    "    NMLikelihoodConfig,\n",
+    ")\n",
+    "from careamics.config.nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig\n",
+    "from careamics.lightning import VAEModule\n",
+    "from careamics.models.lvae.noise_models import noise_model_factory"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Set some parameters for the current training simulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "img_size: int = 64\n",
+    "\"\"\"Spatial size of the input image.\"\"\"\n",
+    "target_channels: int = 2\n",
+    "\"\"\"Number of channels in the target image.\"\"\"\n",
+    "multiscale_count: int = 5\n",
+    "\"\"\"The number of LC inputs plus one (the actual input).\"\"\"\n",
+    "predict_logvar: Optional[Literal[\"pixelwise\"]] = \"pixelwise\"\n",
+    "\"\"\"Whether to compute also the log-variance as LVAE output.\"\"\"\n",
+    "loss_type: Optional[Literal[\"musplit\", \"denoisplit\", \"denoisplit_musplit\"]] = \"musplit\"\n",
+    "\"\"\"The type of reconstruction loss (i.e., likelihood) to use.\"\"\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Create `Dataset` and `Dataloader`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 1.1. Dummy Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DummyDataset(Dataset):\n",
+    "    def __init__(\n",
+    "        self, \n",
+    "        img_size: int = 64, \n",
+    "        target_ch: int = 1,\n",
+    "        multiscale_count: int = 1,\n",
+    "    ):\n",
+    "        self.num_samples = 100\n",
+    "        self.img_size = img_size\n",
+    "        self.target_ch = target_ch\n",
+    "        self.multiscale_count = multiscale_count\n",
+    "    \n",
+    "    def __len__(self):\n",
+    "        return self.num_samples\n",
+    "    \n",
+    "    def __getitem__(self, idx: int):\n",
+    "        input_ = torch.randn(self.multiscale_count, self.img_size, self.img_size)\n",
+    "        target = torch.randn(self.target_ch, self.img_size, self.img_size)\n",
+    "        return input_, target\n",
+    "\n",
+    "def dummy_dataloader(\n",
+    "    batch_size: int = 1,\n",
+    "    img_size: int = 64,\n",
+    "    target_ch: int = 1,\n",
+    "    multiscale_count: int = 1,\n",
+    "):\n",
+    "    dataset = DummyDataset(\n",
+    "        img_size=img_size,\n",
+    "        target_ch=target_ch,\n",
+    "        multiscale_count=multiscale_count,\n",
+    "    )\n",
+    "    return DataLoader(dataset, batch_size=batch_size, num_workers=3, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dloader = dummy_dataloader(\n",
+    "    img_size=img_size,\n",
+    "    target_ch=target_channels,\n",
+    "    multiscale_count=multiscale_count,\n",
+    ")\n",
+    "input_, target = next(iter(dloader))\n",
+    "input_.shape, target.shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 1.2. Real Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Instantiate the lightning module"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_dummy_noise_model(\n",
+    "    save_path: Optional[Union[Path, str]] = None,\n",
+    "    n_gaussians: int = 3,\n",
+    "    n_coeffs: int = 3,\n",
+    ") -> Path:\n",
+    "    weights = np.random.rand(3*n_gaussians, n_coeffs)\n",
+    "    nm_dict = {\n",
+    "        \"trained_weight\": weights,\n",
+    "        \"min_signal\": np.array([0]),\n",
+    "        \"max_signal\": np.array([2**16 - 1]),\n",
+    "        \"min_sigma\": 0.125,\n",
+    "    }\n",
+    "    out_path = Path(save_path) / \"dummy_noise_model.npz\"\n",
+    "    np.savez(out_path, **nm_dict)\n",
+    "    return out_path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_split_lightning_model(\n",
+    "    algorithm: str,\n",
+    "    loss_type: str,\n",
+    "    multiscale_count: int = 1,\n",
+    "    predict_logvar: Optional[Literal[\"pixelwise\"]] = None,\n",
+    "    target_ch: int = 1,\n",
+    "    NM_path: Optional[Path] = None,\n",
+    ") -> VAEModule:\n",
+    "    \"\"\"Instantiate the muSplit lightining model.\"\"\"\n",
+    "    lvae_config = LVAEModel(\n",
+    "        architecture=\"LVAE\",\n",
+    "        input_shape=64,\n",
+    "        multiscale_count=multiscale_count,\n",
+    "        z_dims=[128, 128, 128, 128],\n",
+    "        output_channels=target_ch,\n",
+    "        predict_logvar=predict_logvar,\n",
+    "    )\n",
+    "\n",
+    "    # gaussian likelihood\n",
+    "    if loss_type in [\"musplit\", \"denoisplit_musplit\"]:\n",
+    "        gaussian_lik_config = GaussianLikelihoodConfig(\n",
+    "            predict_logvar=predict_logvar,\n",
+    "            logvar_lowerbound=0.0,\n",
+    "        )\n",
+    "    else:\n",
+    "        gaussian_lik_config = None\n",
+    "    # noise model likelihood\n",
+    "    if loss_type in [\"denoisplit\", \"denoisplit_musplit\"]:\n",
+    "        if NM_path is None:\n",
+    "            NM_path = create_dummy_noise_model(Path(\"./\"), 3, 3)\n",
+    "        gmm = GaussianMixtureNMConfig(\n",
+    "            model_type=\"GaussianMixtureNoiseModel\",\n",
+    "            path=NM_path,\n",
+    "        )\n",
+    "        noise_model_config = MultiChannelNMConfig(noise_models=[gmm] * target_ch)\n",
+    "        nm = noise_model_factory(noise_model_config)\n",
+    "        nm_lik_config = NMLikelihoodConfig(noise_model=nm)\n",
+    "    else:\n",
+    "        noise_model_config = None\n",
+    "        nm_lik_config = None\n",
+    "\n",
+    "    vae_config = VAEAlgorithmConfig(\n",
+    "        algorithm_type=\"vae\",\n",
+    "        algorithm=algorithm,\n",
+    "        loss=loss_type,\n",
+    "        model=lvae_config,\n",
+    "        gaussian_likelihood_model=gaussian_lik_config,\n",
+    "        noise_model=noise_model_config,\n",
+    "        noise_model_likelihood_model=nm_lik_config,\n",
+    "    )\n",
+    "\n",
+    "    return VAEModule(\n",
+    "        algorithm_config=vae_config,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "algo = \"musplit\" if loss_type == \"musplit\" else \"denoisplit\"\n",
+    "lightning_model = create_split_lightning_model(\n",
+    "    algorithm=algo,\n",
+    "    loss_type=loss_type,\n",
+    "    multiscale_count=multiscale_count,\n",
+    "    predict_logvar=predict_logvar,\n",
+    "    target_ch=target_channels,\n",
+    "    NM_path=None\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Set utils for training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datetime import datetime\n",
+    "\n",
+    "from careamics.lvae_training.train_utils import get_new_model_version\n",
+    "\n",
+    "def get_new_model_version(model_dir: Union[Path, str]) -> int:\n",
+    "    \"\"\"Create a unique version ID for a new model run.\"\"\"\n",
+    "    versions = []\n",
+    "    for version_dir in os.listdir(model_dir):\n",
+    "        try:\n",
+    "            versions.append(int(version_dir))\n",
+    "        except:\n",
+    "            print(\n",
+    "                f\"Invalid subdirectory:{model_dir}/{version_dir}. Only integer versions are allowed\"\n",
+    "            )\n",
+    "            exit()\n",
+    "    if len(versions) == 0:\n",
+    "        return \"0\"\n",
+    "    return f\"{max(versions) + 1}\"\n",
+    "\n",
+    "def get_workdir(\n",
+    "    root_dir: str,\n",
+    "    model_name: str,\n",
+    ") -> tuple[Path, Path]:\n",
+    "    \"\"\"Get the workdir for the current model.\n",
+    "    \n",
+    "    It has the following structure: \"root_dir/YYMM/model_name/version\"\n",
+    "    \"\"\"\n",
+    "    rel_path = datetime.now().strftime(\"%y%m\")\n",
+    "    cur_workdir = os.path.join(root_dir, rel_path)\n",
+    "    Path(cur_workdir).mkdir(exist_ok=True)\n",
+    "\n",
+    "    rel_path = os.path.join(rel_path, model_name)\n",
+    "    cur_workdir = os.path.join(root_dir, rel_path)\n",
+    "    Path(cur_workdir).mkdir(exist_ok=True)\n",
+    "\n",
+    "    rel_path = os.path.join(rel_path, get_new_model_version(cur_workdir))\n",
+    "    cur_workdir = os.path.join(root_dir, rel_path)\n",
+    "    try:\n",
+    "        Path(cur_workdir).mkdir(exist_ok=False)\n",
+    "    except FileExistsError:\n",
+    "        print(\n",
+    "            f\"Workdir {cur_workdir} already exists.\"\n",
+    "        )\n",
+    "    return cur_workdir, rel_path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ROOT_DIR = \"/group/jug/federico/careamics_training/refac_v2/\"\n",
+    "workdir, exp_tag = get_workdir(ROOT_DIR, \"dummy_debugging\")\n",
+    "print(f\"Current workdir: {workdir}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the logger\n",
+    "custom_logger = WandbLogger(\n",
+    "    name=os.path.join(socket.gethostname(), exp_tag),\n",
+    "    save_dir=workdir,\n",
+    "    project=\"careamics_debugging_LVAE\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define callbacks (e.g., ModelCheckpoint, EarlyStopping, etc.)\n",
+    "early_stopping_config = EarlyStoppingModel(\n",
+    "    monitor=\"val_loss\",\n",
+    "    min_delta=1e-6,\n",
+    "    patience=10,\n",
+    "    mode=\"min\",\n",
+    "    verbose=True,\n",
+    ")\n",
+    "checkpoint_config = CheckpointModel(\n",
+    "    monitor=\"val_loss\",\n",
+    "    save_top_k=2,\n",
+    "    mode=\"min\",\n",
+    ")\n",
+    "custom_callbacks = [\n",
+    "    EarlyStopping(**early_stopping_config.model_dump()), \n",
+    "    ModelCheckpoint(**checkpoint_config.model_dump()),\n",
+    "    LearningRateMonitor(logging_interval=\"epoch\")\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Save AlgorithmConfig\n",
+    "with open(os.path.join(workdir, \"algorithm_config.json\"), \"w\") as f:\n",
+    "    f.write(lightning_model.algorithm_config.model_dump_json())\n",
+    "\n",
+    "custom_logger.experiment.config.update(\n",
+    "    lightning_model.algorithm_config.model_dump()    \n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Train the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = Trainer(\n",
+    "    max_epochs=10,\n",
+    "    accelerator=\"cpu\",\n",
+    "    enable_progress_bar=True,\n",
+    "    logger=custom_logger,\n",
+    "    callbacks=custom_callbacks,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.fit(\n",
+    "    model=lightning_model,\n",
+    "    train_dataloaders=dloader,\n",
+    "    val_dataloaders=dloader,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "train_lvae",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

careamics-0.0.3/mypy.ini

@@ -0,0 +1,17 @@
+[mypy]
+ignore_missing_imports = True
+
+[mypy-careamics.lvae_training.*]
+follow_imports = skip
+
+[mypy-careamics.models.lvae.*]
+follow_imports = skip
+
+[mypy-careamics.losses.loss_factory]
+follow_imports = skip
+
+[mypy-careamics.losses.lvae.losses]
+follow_imports = skip
+
+[mypy-careamics.config.likelihood_model]
+follow_imports = skip

{careamics-0.0.2 → careamics-0.0.3}/pyproject.toml

@@ -23,8 +23,12 @@ readme = "README.md"
 requires-python = ">=3.9"
 license = { text = "BSD-3-Clause" }
 authors = [
+    { name = 'CAREamics team', email = 'rse@fht.org' },
+    { name = 'Ashesh', email = 'ashesh.ashesh@fht.org' },
+    { name = 'Federico Carrara', email = 'federico.carrara@fht.org' },
     { name = 'Melisande Croft', email = 'melisande.croft@fht.org' },
     { name = 'Joran Deschamps', email = 'joran.deschamps@fht.org' },
+    { name = 'Vera Galinova', email = 'vera.galinova@fht.org' },
     { name = 'Igor Zubarev', email = 'igor.zubarev@fht.org' },
 ]
 classifiers = [

{careamics-0.0.2 → careamics-0.0.3}/src/careamics/careamist.py

@@ -13,10 +13,7 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
-from careamics.config import (
-    Configuration,
-    load_configuration,
-)
+from careamics.config import Configuration, FCNAlgorithmConfig, load_configuration
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -25,7 +22,7 @@ from careamics.config.support import (
 )
 from careamics.dataset.dataset_utils import reshape_array
 from careamics.lightning import (
-    CAREamicsModule,
+    FCNModule,
     HyperParametersCallback,
     PredictDataModule,
     ProgressBarCallback,
@@ -148,9 +145,12 @@ class CAREamist:
             self.cfg = source
 
             # instantiate model
-            self.model = CAREamicsModule(
-                algorithm_config=self.cfg.algorithm_config,
-            )
+            if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                self.model = FCNModule(
+                    algorithm_config=self.cfg.algorithm_config,
+                )
+            else:
+                raise NotImplementedError("Architecture not supported.")
 
         # path to configuration file or model
         else:
@@ -164,9 +164,12 @@ class CAREamist:
                 self.cfg = load_configuration(source)
 
                 # instantiate model
-                self.model = CAREamicsModule(
-                    algorithm_config=self.cfg.algorithm_config,
-                )
+                if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                    self.model = FCNModule(
+                        algorithm_config=self.cfg.algorithm_config,
+                    )  # type: ignore
+                else:
+                    raise NotImplementedError("Architecture not supported.")
 
             # attempt loading a pre-trained model
             else:

{careamics-0.0.2 → careamics-0.0.3}/src/careamics/config/__init__.py

@@ -1,7 +1,8 @@
 """Configuration module."""
 
 __all__ = [
-    "AlgorithmConfig",
+    "FCNAlgorithmConfig",
+    "VAEAlgorithmConfig",
     "DataConfig",
     "Configuration",
     "CheckpointModel",
@@ -15,9 +16,9 @@ __all__ = [
     "register_model",
     "CustomModel",
     "clear_custom_models",
+    "GaussianMixtureNMConfig",
+    "MultiChannelNMConfig",
 ]
-
-from .algorithm_model import AlgorithmConfig
 from .architectures import CustomModel, clear_custom_models, register_model
 from .callback_model import CheckpointModel
 from .configuration_factory import (
@@ -31,5 +32,8 @@ from .configuration_model import (
     save_configuration,
 )
 from .data_model import DataConfig
+from .fcn_algorithm_model import FCNAlgorithmConfig
 from .inference_model import InferenceConfig
+from .nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig
 from .training_model import TrainingConfig
+from .vae_algorithm_model import VAEAlgorithmConfig

{careamics-0.0.2 → careamics-0.0.3}/src/careamics/config/architectures/__init__.py

@@ -4,7 +4,7 @@ __all__ = [
     "ArchitectureModel",
     "CustomModel",
     "UNetModel",
-    "VAEModel",
+    "LVAEModel",
     "clear_custom_models",
     "get_custom_model",
     "register_model",
@@ -12,6 +12,6 @@ __all__ = [
 
 from .architecture_model import ArchitectureModel
 from .custom_model import CustomModel
+from .lvae_model import LVAEModel
 from .register_model import clear_custom_models, get_custom_model, register_model
 from .unet_model import UNetModel
-from .vae_model import VAEModel

{careamics-0.0.2 → careamics-0.0.3}/src/careamics/config/architectures/architecture_model.py

@@ -27,7 +27,7 @@ class ArchitectureModel(BaseModel):
         Returns
         -------
         dict[str, Any]
-            Model as a dictionnary.
+            Model as a dictionary.
         """
         model_dict = super().model_dump(**kwargs)
 

{careamics-0.0.2 → careamics-0.0.3}/src/careamics/config/architectures/custom_model.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import inspect
 from pprint import pformat
 from typing import Any, Literal
 
@@ -23,12 +24,13 @@ class CustomModel(ArchitectureModel):
 
     Attributes
     ----------
-    architecture : Literal["Custom"]
-        Discriminator for the custom model, must be set to "Custom".
+    architecture : Literal["custom"]
+        Discriminator for the custom model, must be set to "custom".
     name : str
         Name of the custom model.
     parameters : CustomParametersModel
-        Parameters of the custom model.
+        All parameters, required for the initialization of the torch module have to be
+        passed here.
 
     Raises
     ------
@@ -57,7 +59,7 @@ class CustomModel(ArchitectureModel):
     ...
     >>> # Create a configuration
     >>> config_dict = {
-    ...     "architecture": "Custom",
+    ...     "architecture": "custom",
     ...     "name": "my_linear",
     ...     "in_features": 10,
     ...     "out_features": 5,
@@ -71,10 +73,9 @@ class CustomModel(ArchitectureModel):
     )
 
     # discriminator used for choosing the pydantic model in Model
-    architecture: Literal["Custom"]
+    architecture: Literal["custom"]
     """Name of the architecture."""
 
-    # name of the custom model
     name: str
     """Name of the custom model."""
 
@@ -120,10 +121,12 @@ class CustomModel(ArchitectureModel):
             get_custom_model(self.name)(**self.model_dump())
         except Exception as e:
             raise ValueError(
-                f"error while passing parameters to the model {e}. Verify that all "
+                f"while passing parameters to the model {e}. Verify that all "
                 f"mandatory parameters are provided, and that either the {e} accepts "
                 f"*args and **kwargs in its __init__() method, or that no additional"
-                f"parameter is provided."
+                f"parameter is provided. Trace: "
+                f"filename: {inspect.trace()[-1].filename}, function: "
+                f"{inspect.trace()[-1].function}, line: {inspect.trace()[-1].lineno}"
             ) from None
 
         return self