diff --git a/src/extensions/score_draw_uml_funcs/__init__.py b/src/extensions/score_draw_uml_funcs/__init__.py
index 067df1ea..a5eb91ca 100644
--- a/src/extensions/score_draw_uml_funcs/__init__.py
+++ b/src/extensions/score_draw_uml_funcs/__init__.py
@@ -361,7 +361,7 @@ def draw_module(
         if not curr_need:
             logger.info(f"{need}: include with id {need_inc} could not be found")
             continue
-        if curr_need["type"] not in ["comp_arc_sta", "mod_view_sta"]:
+        if curr_need["type"] not in ["comp", "mod"]:
             continue
         sub_structure, sub_linkage, proc_impl_interfaces, proc_used_interfaces = (
             draw_comp_incl_impl_int(
diff --git a/src/extensions/score_draw_uml_funcs/helpers.py b/src/extensions/score_draw_uml_funcs/helpers.py
index 37b1eddf..f82e2396 100644
--- a/src/extensions/score_draw_uml_funcs/helpers.py
+++ b/src/extensions/score_draw_uml_funcs/helpers.py
@@ -28,7 +28,7 @@ def gen_format(need: dict[str, str]) -> str:
 
     style = ""
 
-    if "comp_arc_sta" in need["type"] and need["safety"] == "ASIL_B":
+    if "comp" in need["type"] and need["safety"] == "ASIL_B":
         style = "<<asilb>>"
 
     if "real_arc_int" in need["type"]: