Create app.py

app.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import spaces
+
+import gradio as gr
+from threading import Thread
+from transformers import TextIteratorStreamer
+import hashlib
+import os
+
+from transformers import AutoModel, AutoProcessor
+import torch
+
+model = AutoModel.from_pretrained("visheratin/MC-LLaVA-3b", torch_dtype=torch.float16, trust_remote_code=True).to("cuda")
+
+processor = AutoProcessor.from_pretrained("visheratin/MC-LLaVA-3b", trust_remote_code=True)
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+    DTYPE = torch.float16
+else:
+    DEVICE = "cpu"
+    DTYPE = torch.float32
+
+def cached_vision_process(image, max_crops, num_tokens):
+    image_hash = hashlib.sha256(image.tobytes()).hexdigest()
+    cache_path = f"visual_cache/{image_hash}-{max_crops}-{num_tokens}.pt"
+    if os.path.exists(cache_path):
+        return torch.load(cache_path).to(DEVICE, dtype=DTYPE)
+    else:
+        processor_outputs = processor.image_processor([image], max_crops)
+        pixel_values = processor_outputs["pixel_values"]
+        pixel_values = [
+            value.to(model.device).to(model.dtype) for value in pixel_values
+        ]
+        coords = processor_outputs["coords"]
+        coords = [value.to(model.device).to(model.dtype) for value in coords]
+        image_outputs = model.vision_model(pixel_values, coords, num_tokens)
+        image_features = model.multi_modal_projector(image_outputs)
+        os.makedirs("visual_cache", exist_ok=True)
+        torch.save(image_features, cache_path)
+        return image_features.to(DEVICE, dtype=DTYPE)
+
+@spaces.GPU(duration=20)
+def answer_question(image, question, max_crops, num_tokens):
+    prompt = f"""<|im_start|>user
+<image>
+{question}<|im_end|>
+<|im_start|>assistant
+"""
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_special_tokens=True)
+    inputs = processor(prompt, [image], model, max_crops=max_crops, num_tokens=num_tokens)
+    generation_kwargs = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+        "image_features": cached_vision_process(image, max_crops, num_tokens),
+        "streamer": streamer,
+        "max_length": 1000,
+        "use_cache": True,
+        "eos_token_id": processor.tokenizer.eos_token_id,
+        "pad_token_id": processor.tokenizer.eos_token_id,
+    }
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        if len(buffer) > 1:
+            yield buffer
+
+
+with gr.Blocks() as demo:
+    gr.HTML("<h1 class='gradio-heading'><center>MC-LLaVA 3B</center></h1>")
+    gr.HTML(
+        "<center><p class='gradio-sub-heading'>MC-LLaVA 3B is a model that can answer questions about small details in high-resolution images. Check out the <a href='https://huggingface.co/visheratin/MC-LLaVA-3b'>model card</a> for more details. If you have any questions or ideas hot to make the model better, <a href='https://x.com/visheratin'>let me know</a></p></center>"
+    )
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Textbox(
+                label="Question", placeholder="e.g. What is this?", scale=4
+            )
+            submit = gr.Button(
+                "Submit",
+                scale=1,
+            )
+        with gr.Row():
+            max_crops = gr.Slider(minimum=0, maximum=200, step=5, value=0, label="Max crops")
+            num_tokens = gr.Slider(minimum=728, maximum=2184, step=10, value=728, label="Number of image tokens")
+        with gr.Row():
+            img = gr.Image(type="pil", label="Upload or Drag an Image")
+            output = gr.TextArea(label="Answer")
+
+    submit.click(answer_question, [img, prompt, max_crops, num_tokens], output)
+    prompt.submit(answer_question, [img, prompt, max_crops, num_tokens], output)
+
+demo.queue().launch(debug=True)