add 8bit

requirements.txt

@@ -5,3 +5,5 @@ PyMuPDF~=1.22
 peft>=0.2.0
 transformers
 gradio
+accelerate
+bitsandbytes

webui.py

@@ -13,7 +13,7 @@ from typing import Optional
 from PIL import Image
 import fitz
 import gradio as gr
-from transformers import TextIteratorStreamer, pipeline, ImageToTextPipeline
+from transformers import TextIteratorStreamer, pipeline, ImageToTextPipeline, AutoModelForPreTraining, AutoProcessor
 
 from infer import TikzDocument, TikzGenerator
 
@@ -23,11 +23,19 @@ models = {
 }
 
 
+def is_8bit(model_name):
+    return "waleko/TikZ-llava" in model_name
+
+
 @lru_cache(maxsize=1)
 def cached_load(model_name, **kwargs) -> ImageToTextPipeline:
     gr.Info("Instantiating model. Could take a while...") # type: ignore
-    # noinspection PyTypeChecker
-    return pipeline("image-to-text", model=model_name, **kwargs)
+    if not is_8bit(model_name):
+        return pipeline("image-to-text", model=model_name, **kwargs)
+    else:
+        model = AutoModelForPreTraining.from_pretrained(model_name, load_in_8bit=True, **kwargs)
+        processor = AutoProcessor.from_pretrained(model_name)
+        return pipeline(task="image-to-text", model=model, tokenizer=processor.tokenizer, image_processor=processor.image_processor)
 
 
 def convert_to_svg(pdf):