diff --git a/crystal_toolkit/components/diffraction.py b/crystal_toolkit/components/diffraction.py
index 64d74e78..1696106a 100644
--- a/crystal_toolkit/components/diffraction.py
+++ b/crystal_toolkit/components/diffraction.py
@@ -503,8 +503,8 @@ def pattern_from_struct(struct, rad_source):
             Output(self.id("large_cell_note"), "children"),
             Input(self.id("structure"), "data"),
         )
-        def update_message(structure):
-            if len(self.from_data(structure).sites) < SITES_LIMIT:
+        def update_message(structure: Structure | None):
+            if structure is None or len(self.from_data(structure).sites) < SITES_LIMIT:
                 return html.Div([])
             return MessageContainer(
                 MessageBody(