diff --git a/mail_unique_layout/models/__init__.py b/mail_unique_layout/models/__init__.py
index 20a3b8cd1..6d9bf6767 100644
--- a/mail_unique_layout/models/__init__.py
+++ b/mail_unique_layout/models/__init__.py
@@ -1,2 +1,3 @@
 from . import mail_thread
 from . import mail_template
+from . import base
diff --git a/mail_unique_layout/models/base.py b/mail_unique_layout/models/base.py
new file mode 100644
index 000000000..24bb41554
--- /dev/null
+++ b/mail_unique_layout/models/base.py
@@ -0,0 +1,14 @@
+# Copyright 2024 Akretion (https://www.akretion.com).
+# @author Sébastien BEAU <sebastien.beau@akretion.com>
+# License AGPL-3.0 or later (https://www.gnu.org/licenses/agpl).
+
+
+from odoo import api, models
+
+
+class Base(models.AbstractModel):
+    _inherit = "base"
+
+    @api.model
+    def _get_mail_layout(self):
+        return "mail_unique_layout.general_mail_layout"
diff --git a/mail_unique_layout/models/ir_ui_view.py b/mail_unique_layout/models/ir_ui_view.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/mail_unique_layout/models/mail_template.py b/mail_unique_layout/models/mail_template.py
index dd38f8236..0e7753cdf 100644
--- a/mail_unique_layout/models/mail_template.py
+++ b/mail_unique_layout/models/mail_template.py
@@ -17,10 +17,11 @@ def send_mail(
         email_values=None,
         notif_layout=False,
     ):
+        record = self.env[self.model].browse(res_id)
         return super().send_mail(
             res_id,
             force_send=force_send,
             raise_exception=raise_exception,
             email_values=email_values,
-            notif_layout=self.env["mail.thread"]._get_unique_layout_ref(),
+            notif_layout=record._get_mail_layout(),
         )
diff --git a/mail_unique_layout/models/mail_thread.py b/mail_unique_layout/models/mail_thread.py
index 5395e8308..5e90535f6 100644
--- a/mail_unique_layout/models/mail_thread.py
+++ b/mail_unique_layout/models/mail_thread.py
@@ -7,18 +7,7 @@
 class MailThread(models.AbstractModel):
     _inherit = "mail.thread"
 
-    @api.model
-    def _get_default_unique_layout_ref(self):
-        return "mail_unique_layout.general_mail_layout"
-
-    @api.model
-    def _get_unique_layout_ref(self):
-        if self._context.get("force_mail_uniq_layout_id"):
-            return self._context["force_mail_uniq_layout_id"]
-        else:
-            return self._get_default_unique_layout_ref()
-
     @api.returns("mail.message", lambda value: value.id)
     def message_post(self, *kw, **kwargs):
-        kwargs["email_layout_xmlid"] = self._get_unique_layout_ref()
+        kwargs["email_layout_xmlid"] = self._get_mail_layout()
         return super().message_post(*kw, **kwargs)
diff --git a/mail_unique_layout/wizards/mail_compose_message.py b/mail_unique_layout/wizards/mail_compose_message.py
index b26c1b556..0ace7b44c 100644
--- a/mail_unique_layout/wizards/mail_compose_message.py
+++ b/mail_unique_layout/wizards/mail_compose_message.py
@@ -9,7 +9,6 @@ class MailComposeMessage(models.TransientModel):
     _inherit = "mail.compose.message"
 
     def send_mail(self, auto_commit=False):
-        self = self.with_context(
-            custom_layout=self.env["mail.thread"]._get_unique_layout_ref()
-        )
+        record = self.env[self.model].browse(self.res_id)
+        self = self.with_context(custom_layout=record._get_mail_layout())
         return super().send_mail(auto_commit=auto_commit)