feat(mm): llava model loading supports partial loading; fix OOM crash on initial load

The model manager has two types of model cache entries:
- `CachedModelOnlyFullLoad`: The model may only ever be loaded and unloaded as a single object.
- `CachedModelWithPartialLoad`: The model may be partially loaded and unloaded.

Partial loaded is enabled by overwriting certain torch layer classes, adding the ability to autocast the layer to a device on-the-fly. See `CustomLinear` for an example.

So, to take advantage of partial loading and be cached as a `CachedModelWithPartialLoad`, the model must inherit from `torch.nn.Module`.

The LLaVA classes provided by `transformers` do inherit from `torch.nn.Module`, but we wrap those classes in a separate class called `LlavaOnevisionModel`. The wrapper encapsulate both the LLaVA model and its "processor" - a lightweight class that prepares model inputs like text and images.

While it is more elegant to encapsulate both model and processor classes in a single entity, this prevents the model cache from enabling partial loading for the chunky vLLM model.

Fixing this involved a few changes.
- Update the `LlavaOnevisionModelLoader` class to operate on the vLLM model directly, instead the `LlavaOnevisionModel` wrapper class.
- Instantiate the processor directly in the node. The processor is lightweight and does its business on the CPU. We don't need to worry about caching in the model manager.
- Remove caching support code from the `LlavaOnevisionModel` wrapper class. It's not needed, because we do not cache this class. The class now only handles running the models provided to it.
- Rename `LlavaOnevisionModel` to `LlavaOnevisionPipeline` to better represent its purpose.

These changes have a bonus effect of fixing an OOM crash when initially loading the models. This was most apparent when loading LLaVA 7B, which is pretty chunky.

The initial load is onto CPU RAM. In the old version of the loaders, we ignored the loader's target dtype for the initial load. Instead, we loaded the model at `transformers`'s "default" dtype of fp32.

LLaVA 7B is fp16 and weighs ~17GB. Loading as fp32 means we need double that amount (~34GB) of CPU RAM. Many users only have 32GB RAM, so this causes a _CPU_ OOM - which is a hard crash of the whole process.

With the updated loaders, the initial load logic now uses the target dtype for the initial load. LLaVA now needs the expected ~17GB RAM for its initial load.

PS: If we didn't make the accompanying partial loading changes, we still could have solved this OOM. We'd just need to pass the initial load dtype to the wrapper class and have it load on that dtype. But we may as well fix both issues.

PPS: There are other models whose model classes are wrappers around a torch module class, and thus cannot be partially loaded. However, these models are typically fairly small and/or are run only on their own, so they don't benefit as much from partial loading. It's the really big models (like LLaVA 7B) that benefit most from the partial loading.

Author: psychedelicious

invokeai/app/invocations/llava_onevision_vllm.py

@@ -3,13 +3,14 @@
 import torch
 from PIL.Image import Image
 from pydantic import field_validator
+from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration, LlavaOnevisionProcessor
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
 from invokeai.app.invocations.fields import FieldDescriptions, ImageField, InputField, UIComponent, UIType
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.invocations.primitives import StringOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.llava_onevision_model import LlavaOnevisionModel
+from invokeai.backend.llava_onevision_pipeline import LlavaOnevisionPipeline
 from invokeai.backend.util.devices import TorchDevice
 
 
@@ -54,10 +55,17 @@ def _get_images(self, context: InvocationContext) -> list[Image]:
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> StringOutput:
         images = self._get_images(context)
+        model_config = context.models.get_config(self.vllm_model)
 
-        with context.models.load(self.vllm_model) as vllm_model:
-            assert isinstance(vllm_model, LlavaOnevisionModel)
-            output = vllm_model.run(
+        with context.models.load(self.vllm_model).model_on_device() as (_, model):
+            assert isinstance(model, LlavaOnevisionForConditionalGeneration)
+
+            model_abs_path = context.models.get_absolute_path(model_config)
+            processor = AutoProcessor.from_pretrained(model_abs_path, local_files_only=True)
+            assert isinstance(processor, LlavaOnevisionProcessor)
+
+            model = LlavaOnevisionPipeline(model, processor)
+            output = model.run(
                 prompt=self.prompt,
                 images=images,
                 device=TorchDevice.choose_torch_device(),

invokeai/backend/llava_onevision_model.py renamed to invokeai/backend/llava_onevision_pipeline.py

@@ -1,26 +1,15 @@
-from pathlib import Path
-from typing import Optional
-
 import torch
 from PIL.Image import Image
-from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration, LlavaOnevisionProcessor
+from transformers import LlavaOnevisionForConditionalGeneration, LlavaOnevisionProcessor
 
-from invokeai.backend.raw_model import RawModel
 
+class LlavaOnevisionPipeline:
+    """A wrapper for a LLaVA Onevision model + processor."""
 
-class LlavaOnevisionModel(RawModel):
     def __init__(self, vllm_model: LlavaOnevisionForConditionalGeneration, processor: LlavaOnevisionProcessor):
         self._vllm_model = vllm_model
         self._processor = processor
 
-    @classmethod
-    def load_from_path(cls, path: str | Path):
-        vllm_model = LlavaOnevisionForConditionalGeneration.from_pretrained(path, local_files_only=True)
-        assert isinstance(vllm_model, LlavaOnevisionForConditionalGeneration)
-        processor = AutoProcessor.from_pretrained(path, local_files_only=True)
-        assert isinstance(processor, LlavaOnevisionProcessor)
-        return cls(vllm_model, processor)
-
     def run(self, prompt: str, images: list[Image], device: torch.device, dtype: torch.dtype) -> str:
         # TODO(ryand): Tune the max number of images that are useful for the model.
         if len(images) > 3:
@@ -44,13 +33,3 @@ def run(self, prompt: str, images: list[Image], device: torch.device, dtype: tor
         # The output_str will include the prompt, so we extract the response.
         response = output_str.split("assistant\n", 1)[1].strip()
         return response
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        self._vllm_model.to(device=device, dtype=dtype)
-
-    def calc_size(self) -> int:
-        """Get size of the model in memory in bytes."""
-        # HACK(ryand): Fix this issue with circular imports.
-        from invokeai.backend.model_manager.load.model_util import calc_module_size
-
-        return calc_module_size(self._vllm_model)

invokeai/backend/model_manager/load/model_loaders/llava_onevision.py

@@ -1,7 +1,8 @@
 from pathlib import Path
 from typing import Optional
 
-from invokeai.backend.llava_onevision_model import LlavaOnevisionModel
+from transformers import LlavaOnevisionForConditionalGeneration
+
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
 )
@@ -23,6 +24,8 @@ def _load_model(
             raise ValueError("Unexpected submodel requested for LLaVA OneVision model.")
 
         model_path = Path(config.path)
-        model = LlavaOnevisionModel.load_from_path(model_path)
-        model.to(dtype=self._torch_dtype)
+        model = LlavaOnevisionForConditionalGeneration.from_pretrained(
+            model_path, local_files_only=True, torch_dtype=self._torch_dtype
+        )
+        assert isinstance(model, LlavaOnevisionForConditionalGeneration)
         return model