diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py
index 1cc0ca620..2f89863ae 100644
--- a/src/backend/langflow/interface/custom/custom_component.py
+++ b/src/backend/langflow/interface/custom/custom_component.py
@@ -5,7 +5,6 @@ from langflow.interface.custom.component import Component
 
 from langflow.utils import validate
 
-from uuid import UUID
 from langflow.database.base import session_getter
 from langflow.database.models.flow import Flow
 from pydantic import Extra
@@ -44,13 +43,15 @@ class CustomComponent(Component, extra=Extra.allow):
         return True
 
     def is_check_valid(self) -> bool:
-        return self._class_template_validation(self.code)
+        return self._class_template_validation(self.code) if self.code else False
 
     def get_code_tree(self, code: str):
         return super().get_code_tree(code)
 
     @property
     def get_function_entrypoint_args(self) -> str:
+        if not self.code:
+            return ""
         tree = self.get_code_tree(self.code)
 
         component_classes = [
@@ -78,6 +79,8 @@ class CustomComponent(Component, extra=Extra.allow):
 
     @property
     def get_function_entrypoint_return_type(self) -> str:
+        if not self.code:
+            return ""
         tree = self.get_code_tree(self.code)
 
         component_classes = [
@@ -138,7 +141,7 @@ class CustomComponent(Component, extra=Extra.allow):
     def get_function(self):
         return validate.create_function(self.code, self.function_entrypoint_name)
 
-    def load_flow(self, flow_id: UUID = None):
+    def load_flow(self, flow_id: str):
         from langflow.processing.process import build_sorted_vertices_with_caching
 
         with session_getter() as session:
diff --git a/src/backend/langflow/utils/util.py b/src/backend/langflow/utils/util.py
index ce5f03bdf..f68c9dbe2 100644
--- a/src/backend/langflow/utils/util.py
+++ b/src/backend/langflow/utils/util.py
@@ -9,7 +9,7 @@ from docstring_parser import parse  # type: ignore
 from langflow.template.frontend_node.constants import FORCE_SHOW_FIELDS
 from langflow.utils import constants
 from langflow.utils.logger import logger
-from multiprocess import cpu_count
+from multiprocess import cpu_count  # type: ignore
 
 
 def build_template_from_function(
@@ -301,13 +301,15 @@ def get_type(value: Any) -> Union[str, type]:
     return _type if isinstance(_type, str) else _type.__name__
 
 
-def remove_optional_wrapper(_type: str) -> str:
+def remove_optional_wrapper(_type: Union[str, type]) -> str:
     """
     Removes the 'Optional' wrapper from the type string.
 
     Returns:
         The type string with the 'Optional' wrapper removed.
     """
+    if isinstance(_type, type):
+        _type = str(_type)
     if "Optional" in _type:
         _type = _type.replace("Optional[", "")[:-1]