From b84e3260cb04c7230804bad0b2388e3bb6897b29 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 29 Oct 2023 12:50:54 -0600 Subject: [PATCH] Fix llava import --- .../dataset_tools/tools/llava_utils.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/extensions_built_in/dataset_tools/tools/llava_utils.py b/extensions_built_in/dataset_tools/tools/llava_utils.py index 6619da9..24455a3 100644 --- a/extensions_built_in/dataset_tools/tools/llava_utils.py +++ b/extensions_built_in/dataset_tools/tools/llava_utils.py @@ -1,9 +1,3 @@ -try: - from llava.model import LlavaLlamaForCausalLM -except ImportError: - # print("You need to manually install llava -> pip install --no-deps git+https://github.com/haotian-liu/LLaVA.git") - print("You need to manually install llava -> pip install --no-deps git+https://github.com/haotian-liu/LLaVA.git") - raise from .caption import default_long_prompt, default_short_prompt, default_replacements, clean_caption @@ -20,6 +14,13 @@ img_ext = ['.jpg', '.jpeg', '.png', '.webp'] class LLaVAImageProcessor: def __init__(self, device='cuda'): + try: + from llava.model import LlavaLlamaForCausalLM + except ImportError: + # print("You need to manually install llava -> pip install --no-deps git+https://github.com/haotian-liu/LLaVA.git") + print( + "You need to manually install llava -> pip install --no-deps git+https://github.com/haotian-liu/LLaVA.git") + raise self.device = device self.model: LlavaLlamaForCausalLM = None self.tokenizer: AutoTokenizer = None