fix selector

webui.py

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