add new model

webui.py

@@ -25,8 +25,9 @@ from infer import TikzDocument, TikzGenerator
 
 # assets = files(__package__) / "assets" if __package__ else files("assets") / "."
 models = {
-    "pix2tikz": "pix2tikz/mixed_e362_step201.pth",
-    "llava-1.5-7b-hf": "waleko/TikZ-llava-1.5-7b"
+    "pix2tikz": {"model": "pix2tikz/mixed_e362_step201.pth"},
+    "llava-1.5-7b-hf": {"model": "waleko/TikZ-llava-1.5-7b"},
+    "new llava-1.5-7b-hf": {"model": "waleko/TikZ-llava-1.5-7b", "revision": "v2"},
 }
 
 
@@ -35,12 +36,13 @@ def is_quantization(model_name):
 
 
 @lru_cache(maxsize=1)
-def cached_load(model_name, **kwargs) -> ImageToTextPipeline:
+def cached_load(model_dict, **kwargs) -> ImageToTextPipeline:
+    model_name = model_dict["model"]
     gr.Info("Instantiating model. Could take a while...") # type: ignore
     if not is_quantization(model_name):
-        return pipeline("image-to-text", model=model_name, **kwargs)
+        return pipeline("image-to-text", **model_dict, **kwargs)
     else:
-        model = AutoModelForPreTraining.from_pretrained(model_name, load_in_4bit=True, **kwargs)
+        model = AutoModelForPreTraining.from_pretrained(model_name, load_in_4bit=True, revision=model_dict.get("revision", "main"), **kwargs)
         processor = AutoProcessor.from_pretrained(model_name)
         return pipeline(task="image-to-text", model=model, tokenizer=processor.tokenizer, image_processor=processor.image_processor)
 
@@ -76,7 +78,7 @@ def pix2tikz(
 
 
 def inference(
-    model_name: str,
+    model_dict: dict,
     image_dict: dict,
     temperature: float,
     top_p: float,
@@ -85,12 +87,13 @@ def inference(
 ):
     try:
         image = image_dict['composite']
+        model_name = model_dict["model"]
         if "pix2tikz" in model_name:
             yield pix2tikz(model_name, image, temperature, top_p, top_k, expand_to_square)
             return
 
         generate = TikzGenerator(
-            cached_load(model_name, device_map="auto"),
+            cached_load(model_dict, device_map="auto"),
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,