fix generation

infer.py

@@ -196,34 +196,23 @@ class TikzGenerator:
         self.expand_to_square = expand_to_square
         self.clean_up_output = clean_up_output
         self.pipeline = pipe
-        # self.pipeline.model = torch.compile(self.pipeline.model)
 
         self.default_kwargs = dict(
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            num_return_sequences=1,
-            # max_length=self.pipeline.tokenizer.model_max_length, # type: ignore
             do_sample=True,
-            return_full_text=False,
-            streamer=TextStreamer(self.pipeline.tokenizer, # type: ignore
-                skip_prompt=True,
-                skip_special_tokens=True
-            ),
             max_new_tokens=1024,
         )
 
-        if not stream:
-            self.default_kwargs.pop("streamer")
+        # if not stream:
+        #     self.default_kwargs.pop("streamer")
 
     def generate(self, image: Image.Image, **generate_kwargs):
         prompt = "Assistant helps to write down the TikZ code for the user's image. USER: <image>\nWrite down the TikZ code to draw the diagram shown in the lol. ASSISTANT:"
         tokenizer = self.pipeline.tokenizer
-        print('starting generation')
         text = self.pipeline(image, prompt=prompt, generate_kwargs=(self.default_kwargs | generate_kwargs))[0]["generated_text"] # type: ignore
 
-        print('text generated: ', text)  # TODO: remove
-
         if self.clean_up_output:
             for token in reversed(tokenizer.tokenize(prompt)): # type: ignore
                 # remove leading characters because skip_special_tokens in pipeline
@@ -240,8 +229,6 @@ class TikzGenerator:
             for artifact, replacement in artifacts.items():
                 text = sub(artifact, replacement, text) # type: ignore
 
-        print('cleaned text: ', text)
-
         return TikzDocument(text)
 
 

webui.py

@@ -15,22 +15,22 @@ import fitz
 import gradio as gr
 from transformers import TextIteratorStreamer, pipeline, ImageToTextPipeline, AutoModelForPreTraining, AutoProcessor
 
-from infer import TikzDocument, TikzGenerator
+from .infer import TikzDocument, TikzGenerator
 
 # assets = files(__package__) / "assets" if __package__ else files("assets") / "."
 models = {
-    "Fine-tuned Llava": "waleko/TikZ-llava-1.5-7b"
+    "llava-1.5-7b-hf": "waleko/TikZ-llava-1.5-7b"
 }
 
 
-def is_8bit(model_name):
+def is_quantization(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
-    if not is_8bit(model_name):
+    if not is_quantization(model_name):
         return pipeline("image-to-text", model=model_name, **kwargs)
     else:
         model = AutoModelForPreTraining.from_pretrained(model_name, load_in_8bit=True, **kwargs)
@@ -45,33 +45,35 @@ def convert_to_svg(pdf):
 
 def inference(
     model_name: str,
-    image: Image.Image,
+    image_dict: dict,
     temperature: float,
     top_p: float,
     top_k: int,
     expand_to_square: bool,
 ):
-    generate = TikzGenerator(
-        cached_load(model_name, device_map="auto"),
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        expand_to_square=expand_to_square,
-    )
-    streamer = TextIteratorStreamer(
-        generate.pipeline.tokenizer, # type: ignore
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
-
-    thread = ThreadPool(processes=1)
-    async_result = thread.apply_async(generate, kwds=dict(image=image, streamer=streamer))
-
-    generated_text = ""
-    for new_text in streamer:
-        generated_text += new_text
-        yield generated_text, None, False
-    yield async_result.get().code, None, True
+    try:
+        generate = TikzGenerator(
+            cached_load(model_name, device_map="auto"),
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            expand_to_square=expand_to_square,
+        )
+        streamer = TextIteratorStreamer(
+            generate.pipeline.tokenizer, # type: ignore
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+
+        thread = ThreadPool(processes=1)
+        async_result = thread.apply_async(generate, kwds=dict(image=image_dict['composite'], streamer=streamer))
+        generated_text = ""
+        for new_text in streamer:
+            generated_text += new_text
+            yield generated_text, None, False
+        yield async_result.get().code, None, True
+    except Exception as e:
+        raise gr.Error(f"Internal Error! {e}")
 
 def tex_compile(
     code: str,
@@ -85,7 +87,8 @@ def tex_compile(
         else:
             gr.Warning("TikZ code compiled to an empty image!") # type: ignore
     elif tikzdoc.compiled_with_errors:
-        gr.Warning("TikZ code compiled with errors!") # type: ignore
+        # gr.Warning("TikZ code compiled with errors!") # type: ignore
+        print("TikZ code compiled with errors!")
 
     if rasterize:
         yield tikzdoc.rasterize()
@@ -123,16 +126,15 @@ def remove_darkness(stylable):
     """
     Patch gradio to only contain light mode colors.
     """
-    pass # TODO: remove dark mode colors from the theme
-    # if isinstance(stylable, gr.themes.Base): # remove dark variants from the entire theme
-    #     params = signature(stylable.set).parameters
-    #     colors = {color: getattr(stylable, color.removesuffix("_dark")) for color in dir(stylable) if color in params}
-    #     return stylable.set(**colors)
-    # elif isinstance(stylable, gr.Blocks): # also handle components which do not use the theme (e.g. modals)
-    #     stylable.load(_js="() => document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'))")
-    #     return stylable
-    # else:
-    #     raise ValueError
+    if isinstance(stylable, gr.themes.Base): # remove dark variants from the entire theme
+        params = signature(stylable.set).parameters
+        colors = {color: getattr(stylable, color.removesuffix("_dark")) for color in dir(stylable) if color in params}
+        return stylable.set(**colors)
+    elif isinstance(stylable, gr.Blocks): # also handle components which do not use the theme (e.g. modals)
+        stylable.load(js="() => document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'))")
+        return stylable
+    else:
+        raise ValueError
 
 def build_ui(model=list(models)[0], lock=False, rasterize=False, force_light=False, lock_reason="locked", timeout=120):
     theme = remove_darkness(gr.themes.Soft()) if force_light else gr.themes.Soft()
@@ -148,7 +150,7 @@ def build_ui(model=list(models)[0], lock=False, rasterize=False, force_light=Fal
                 )
                 # caption = gr.Textbox(label="Caption", info=info, placeholder="Type a caption...")
                 # image = gr.Image(label="Image Input", type="pil")
-                image = gr.ImageEditor(label="Image Input", type="pil")
+                image = gr.ImageEditor(label="Image Input", type="pil", sources=['upload', 'clipboard'], value=Image.new('RGB', (336, 336), (255, 255, 255)))
                 label = "Model" + (f" ({lock_reason})" if lock else "")
                 model = gr.Dropdown(label=label, choices=list(models.items()), value=models[model], interactive=not lock) # type: ignore
                 with gr.Accordion(label="Advanced Options", open=False):
@@ -168,7 +170,7 @@ def build_ui(model=list(models)[0], lock=False, rasterize=False, force_light=Fal
                     with gr.TabItem(label:="Compiled Image", id=1):
                         result_image = gr.Image(label=label, show_label=False, show_share_button=rasterize)
                     clear_btn.add([tikz_code, result_image])
-        # TODO: gr.Examples(examples=str(assets), inputs=[image, tikz_code, result_image])
+        gr.Examples(examples=[["https://waleko.github.io/data/image.jpg"]], inputs=[image])
 
         events = list()
         finished = gr.Textbox(visible=False) # hack to cancel compile on canceled inference