Model: Add unload and error messages for vision

If vision is enabled and the model doesn't support it, send an error asking the user to reload. Also, add a method to unload the vision tower. Signed-off-by: kingbri <[email protected]>
theroyallab · Nov 22, 2024 · eadc71a · eadc71a
1 parent c49047e
commit eadc71a
Showing 1 changed file with 17 additions and 3 deletions.
diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
@@ -151,8 +151,14 @@ async def create(cls, model_directory: pathlib.Path, quiet=False, **kwargs):
         # Apply a model's config overrides while respecting user settings
         kwargs = await self.set_model_overrides(**kwargs)
 
-        # Set vision state
+        # Set vision state and error if vision isn't supported on the current model
         self.use_vision = unwrap(kwargs.get("vision"), False)
+        if self.use_vision and not self.config.vision_model_type:
+            raise ValueError(
+                "The provided model does not have vision capabilities that are "
+                "supported by ExllamaV2. "
+                "Please reload with vision disabled."
+            )
 
         # Prepare the draft model config if necessary
         draft_args = unwrap(kwargs.get("draft_model"), {})
@@ -373,8 +379,6 @@ async def create(cls, model_directory: pathlib.Path, quiet=False, **kwargs):
                 self.draft_config.max_input_len = chunk_size
                 self.draft_config.max_attention_size = chunk_size**2
 
-        self.prompt_template = None
-
         # Return the created instance
         return self
 
@@ -848,6 +852,16 @@ async def unload(self, loras_only: bool = False, **kwargs):
                     self.model.unload()
                 self.model = None
 
+                if self.vision_model:
+                    # TODO: Remove this with newer exl2 versions
+                    # Required otherwise unload function won't finish
+                    try:
+                        self.vision_model.unload()
+                    except AttributeError:
+                        pass
+
+                self.vision_model = None
+
                 if self.draft_model:
                     self.draft_model.unload()
                 self.draft_model = None