fix(nodes): ensure SD2 models are pickable in loader/cnet nodes

Author: psychedelicious

invokeai/app/invocations/controlnet.py

@@ -63,14 +63,16 @@ class ControlOutput(BaseInvocationOutput):
     control: ControlField = OutputField(description=FieldDescriptions.control)
 
 
-@invocation("controlnet", title="ControlNet - SD1.5, SDXL", tags=["controlnet"], category="controlnet", version="1.1.3")
+@invocation(
+    "controlnet", title="ControlNet - SD1.5, SD2, SDXL", tags=["controlnet"], category="controlnet", version="1.1.3"
+)
 class ControlNetInvocation(BaseInvocation):
     """Collects ControlNet info to pass to other nodes"""
 
     image: ImageField = InputField(description="The control image")
     control_model: ModelIdentifierField = InputField(
         description=FieldDescriptions.controlnet_model,
-        ui_model_base=[BaseModelType.StableDiffusion1, BaseModelType.StableDiffusionXL],
+        ui_model_base=[BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2, BaseModelType.StableDiffusionXL],
         ui_model_type=ModelType.ControlNet,
     )
     control_weight: Union[float, List[float]] = InputField(

invokeai/app/invocations/model.py

@@ -145,7 +145,7 @@ def invoke(self, context: InvocationContext) -> ModelIdentifierOutput:
 
 @invocation(
     "main_model_loader",
-    title="Main Model - SD1.5",
+    title="Main Model - SD1.5, SD2",
     tags=["model"],
     category="model",
     version="1.0.4",
@@ -155,7 +155,7 @@ class MainModelLoaderInvocation(BaseInvocation):
 
     model: ModelIdentifierField = InputField(
         description=FieldDescriptions.main_model,
-        ui_model_base=BaseModelType.StableDiffusion1,
+        ui_model_base=[BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2],
         ui_model_type=ModelType.Main,
     )
     # TODO: precision?
@@ -485,7 +485,11 @@ def invoke(self, context: InvocationContext) -> SDXLLoRALoaderOutput:
 
 
 @invocation(
-    "vae_loader", title="VAE Model - SD1.5, SDXL, SD3, FLUX", tags=["vae", "model"], category="model", version="1.0.4"
+    "vae_loader",
+    title="VAE Model - SD1.5, SD2, SDXL, SD3, FLUX",
+    tags=["vae", "model"],
+    category="model",
+    version="1.0.4",
 )
 class VAELoaderInvocation(BaseInvocation):
     """Loads a VAE model, outputting a VaeLoaderOutput"""
@@ -495,6 +499,7 @@ class VAELoaderInvocation(BaseInvocation):
         title="VAE",
         ui_model_base=[
             BaseModelType.StableDiffusion1,
+            BaseModelType.StableDiffusion2,
             BaseModelType.StableDiffusionXL,
             BaseModelType.StableDiffusion3,
             BaseModelType.Flux,