[PR-556]: Adding usage of subclasses of ModelClass in Model Upload (#566)

* [PR-556]: Adding usage of subclasses of ModelClass in Model Upload

* [PR-556]: Adding usage of subclasses of ModelClass in Model Upload

Author: srikanthbachala20

clarifai/runners/models/model_builder.py

@@ -43,6 +43,19 @@ def _clear_line(n: int = 1) -> None:
     print(LINE_UP, end=LINE_CLEAR, flush=True)
 
 
+def is_related(object_class, main_class):
+  # Check if the object_class is a subclass of main_class
+  if issubclass(object_class, main_class):
+    return True
+
+  # Check if the object_class is a subclass of any of the parent classes of main_class
+  parent_classes = object_class.__bases__
+  for parent in parent_classes:
+    if main_class in parent.__bases__:
+      return True
+  return False
+
+
 class ModelBuilder:
   DEFAULT_CHECKPOINT_SIZE = 50 * 1024**3  # 50 GiB
 
@@ -125,7 +138,7 @@ def custom_import(name, globals=None, locals=None, fromlist=(), level=0):
     # Find all classes in the model.py file that are subclasses of ModelClass
     classes = [
         cls for _, cls in inspect.getmembers(module, inspect.isclass)
-        if issubclass(cls, ModelClass) and cls.__module__ == module.__name__
+        if is_related(cls, ModelClass) and cls.__module__ == module.__name__
     ]
     #  Ensure there is exactly one subclass of BaseRunner in the model.py file
     if len(classes) != 1: