refactor(nodes): model identifiers

- All models are identified by a key and optionally a submodel type via new model `ModelField`. Previously, a few model types had their own class, but not all of them. This inconsistency just added complexity without any benefit.
- Update all invocation to use the new format.
- In the node API, models are loaded by key or an instance of `ModelField` as a convenience.
- Add an enriched model schema for metadata. It includes key, hash, name, base and type.

invokeai/app/invocations/latent.py

@@ -26,6 +26,7 @@ from diffusers.schedulers import SchedulerMixin as Scheduler
 from PIL import Image, ImageFilter
 from pydantic import field_validator
 from torchvision.transforms.functional import resize as tv_resize
+from transformers import CLIPVisionModelWithProjection
 
 from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
 from invokeai.app.invocations.fields import (
@@ -75,7 +76,7 @@ from .baseinvocation import (
     invocation_output,
 )
 from .controlnet_image_processors import ControlField
-from .model import ModelInfo, UNetField, VaeField
+from .model import ModelField, UNetField, VaeField
 
 if choose_torch_device() == torch.device("mps"):
     from torch import mps
@@ -153,7 +154,7 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
         )
 
         if image_tensor is not None:
-            vae_info = context.models.load(**self.vae.vae.model_dump())
+            vae_info = context.models.load(self.vae.vae)
 
             img_mask = tv_resize(mask, image_tensor.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
             masked_image = image_tensor * torch.where(img_mask < 0.5, 0.0, 1.0)
@@ -244,12 +245,12 @@ class CreateGradientMaskInvocation(BaseInvocation):
 
 def get_scheduler(
     context: InvocationContext,
-    scheduler_info: ModelInfo,
+    scheduler_info: ModelField,
     scheduler_name: str,
     seed: int,
 ) -> Scheduler:
     scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP["ddim"])
-    orig_scheduler_info = context.models.load(**scheduler_info.model_dump())
+    orig_scheduler_info = context.models.load(scheduler_info)
     with orig_scheduler_info as orig_scheduler:
         scheduler_config = orig_scheduler.config
 
@@ -461,7 +462,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         #        and if weight is None, populate with default 1.0?
         controlnet_data = []
         for control_info in control_list:
-            control_model = exit_stack.enter_context(context.models.load(key=control_info.control_model.key))
+            control_model = exit_stack.enter_context(context.models.load(control_info.control_model))
 
             # control_models.append(control_model)
             control_image_field = control_info.image
@@ -523,11 +524,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
         conditioning_data.ip_adapter_conditioning = []
         for single_ip_adapter in ip_adapter:
             ip_adapter_model: Union[IPAdapter, IPAdapterPlus] = exit_stack.enter_context(
-                context.models.load(key=single_ip_adapter.ip_adapter_model.key)
+                context.models.load(single_ip_adapter.ip_adapter_model)
             )
 
-            image_encoder_model_info = context.models.load(key=single_ip_adapter.image_encoder_model.key)
-
+            image_encoder_model_info = context.models.load(single_ip_adapter.image_encoder_model)
             # `single_ip_adapter.image` could be a list or a single ImageField. Normalize to a list here.
             single_ipa_image_fields = single_ip_adapter.image
             if not isinstance(single_ipa_image_fields, list):
@@ -538,6 +538,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             # TODO(ryand): With some effort, the step of running the CLIP Vision encoder could be done before any other
             # models are needed in memory. This would help to reduce peak memory utilization in low-memory environments.
             with image_encoder_model_info as image_encoder_model:
+                assert isinstance(image_encoder_model, CLIPVisionModelWithProjection)
                 # Get image embeddings from CLIP and ImageProjModel.
                 image_prompt_embeds, uncond_image_prompt_embeds = ip_adapter_model.get_image_embeds(
                     single_ipa_images, image_encoder_model
@@ -577,8 +578,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         t2i_adapter_data = []
         for t2i_adapter_field in t2i_adapter:
-            t2i_adapter_model_config = context.models.get_config(key=t2i_adapter_field.t2i_adapter_model.key)
-            t2i_adapter_loaded_model = context.models.load(key=t2i_adapter_field.t2i_adapter_model.key)
+            t2i_adapter_model_config = context.models.get_config(t2i_adapter_field.t2i_adapter_model.key)
+            t2i_adapter_loaded_model = context.models.load(t2i_adapter_field.t2i_adapter_model)
             image = context.images.get_pil(t2i_adapter_field.image.image_name)
 
             # The max_unet_downscale is the maximum amount that the UNet model downscales the latent image internally.
@@ -731,12 +732,13 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
             def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
                 for lora in self.unet.loras:
-                    lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
+                    lora_info = context.models.load(lora.lora)
+                    assert isinstance(lora_info.model, LoRAModelRaw)
                     yield (lora_info.model, lora.weight)
                     del lora_info
                 return
 
-            unet_info = context.models.load(**self.unet.unet.model_dump())
+            unet_info = context.models.load(self.unet.unet)
             assert isinstance(unet_info.model, UNet2DConditionModel)
             with (
                 ExitStack() as exit_stack,
@@ -841,8 +843,8 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     def invoke(self, context: InvocationContext) -> ImageOutput:
         latents = context.tensors.load(self.latents.latents_name)
 
-        vae_info = context.models.load(**self.vae.vae.model_dump())
-
+        vae_info = context.models.load(self.vae.vae)
+        assert isinstance(vae_info.model, (UNet2DConditionModel, AutoencoderKL))
         with set_seamless(vae_info.model, self.vae.seamless_axes), vae_info as vae:
             assert isinstance(vae, torch.nn.Module)
             latents = latents.to(vae.device)
@@ -1064,7 +1066,7 @@ class ImageToLatentsInvocation(BaseInvocation):
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         image = context.images.get_pil(self.image.image_name)
 
-        vae_info = context.models.load(**self.vae.vae.model_dump())
+        vae_info = context.models.load(self.vae.vae)
 
         image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
         if image_tensor.dim() == 3: