Type-safe model purpose for model inputs (#2358)

chaiNNer-org · Nov 30, 2023 · cb86150 · cb86150
1 parent 9a6b394
commit cb86150
Showing 1 changed file with 22 additions and 22 deletions.
diff --git a/backend/src/nodes/properties/inputs/pytorch_inputs.py b/backend/src/nodes/properties/inputs/pytorch_inputs.py
@@ -6,6 +6,7 @@
         ImageModelDescriptor,
         MaskedImageModelDescriptor,
         ModelDescriptor,
+        Purpose,
     )
 except Exception:
     torch = None
@@ -16,6 +17,11 @@
 from api import BaseInput
 
 
+def _model_with_purpose(purpose: set[Purpose]):
+    sub_type = " | ".join('"' + p + '"' for p in purpose)
+    return "PyTorchModel { subType: " + sub_type + " }"
+
+
 class ModelInput(BaseInput):
     """Input a loaded model"""
 
@@ -42,75 +48,69 @@ def __init__(
         label: str = "Model",
         input_type: navi.ExpressionJson = "PyTorchModel",
     ):
+        self.purpose: set[Purpose] = {"SR", "Restoration"}
+
         super().__init__(
             label,
-            navi.intersect(
-                input_type,
-                """PyTorchModel { subType: "SR" | "Restoration" }""",
-            ),
+            navi.intersect(input_type, _model_with_purpose(self.purpose)),
         )
         if torch is not None:
             self.associated_type = ImageModelDescriptor
 
     def enforce(self, value: ModelDescriptor):
         if torch is not None:
+            assert (
+                value.purpose in self.purpose
+            ), "Expected a Super-Resolution or Restoration model."
             assert isinstance(
                 value, ImageModelDescriptor
             ), "Expected a supported single image PyTorch model."
-            assert value.purpose in (
-                "SR",
-                "Restoration",
-            ), "Expected a Super-Resolution or Restoration model."
         return value
 
 
 class FaceModelInput(ModelInput):
     def __init__(
         self, label: str = "Model", input_type: navi.ExpressionJson = "PyTorchModel"
     ):
+        self.purpose: set[Purpose] = {"FaceSR"}
+
         super().__init__(
             label,
-            navi.intersect(
-                input_type,
-                """PyTorchModel { subType: "FaceSR" }""",
-            ),
+            navi.intersect(input_type, _model_with_purpose(self.purpose)),
         )
         if torch is not None:
             self.associated_type = ImageModelDescriptor
 
     def enforce(self, value: ModelDescriptor):
         if torch is not None:
+            assert (
+                value.purpose in self.purpose
+            ), "Expected a Face Super-Resolution model."
             assert isinstance(
                 value, ImageModelDescriptor
             ), "Expected a supported single image PyTorch model."
-            assert value.purpose in (
-                "FaceSR"
-            ), "Expected a Face Super-Resolution model."
         return value
 
 
 class InpaintModelInput(ModelInput):
     def __init__(
         self, label: str = "Model", input_type: navi.ExpressionJson = "PyTorchModel"
     ):
+        self.purpose: set[Purpose] = {"Inpaint"}
+
         super().__init__(
             label,
-            navi.intersect(
-                input_type,
-                """PyTorchModel { subType: "Inpaint" }""",
-            ),
+            navi.intersect(input_type, _model_with_purpose(self.purpose)),
         )
         if torch is not None:
             self.associated_type = MaskedImageModelDescriptor
 
     def enforce(self, value: ModelDescriptor):
         if torch is not None:
+            assert value.purpose in self.purpose, "Expected an Inpainting model."
             assert isinstance(
                 value, MaskedImageModelDescriptor
             ), "Expected a supported masked-image PyTorch model."
-            assert value.purpose in (
-                "Inpaint"
-            ), "Expected a Face Super-Resolution model."
         return value