From eadc71a4c3e889c28d8466035e48ac41f00e87c5 Mon Sep 17 00:00:00 2001 From: kingbri Date: Fri, 22 Nov 2024 14:25:03 -0500 Subject: [PATCH] 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 --- backends/exllamav2/model.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index beb4512..ff11531 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -151,8 +151,14 @@ class ExllamaV2Container: # 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 @@ class ExllamaV2Container: 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 @@ class ExllamaV2Container: 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