Update README.md

README.md

@@ -11,26 +11,22 @@ TiO is an Interactive Visual Grounding Model for Disambiguation.  (WIP)
 
 ## Online / Offline Demo
 
+<img src="https://cdn-uploads.huggingface.co/production/uploads/63ea2c4ed235e6788af30eaa/NrF9poZzUVz1oTZ0uQrun.png" width="300px" />
+
 - [Colab Online Demo](https://colab.research.google.com/drive/195eDITKi6dahnVz8Cum91sNUCF_lFle8?usp=sharing) - Free T4 is available on Google Colab.
 - Gradio Offline Demo:
 
 ```python
-import os; os.system("pip3 install transformers accelerate bitsandbytes gradio fire")
+import os; os.system("pip3 install transformers gradio fire accelerate bitsandbytes > /dev/null")
 from transformers import AutoModel, AutoTokenizer, AutoImageProcessor
+import torch
 
 model_id = "jxu124/TiO"
-model = AutoModel.from_pretrained(
-    model_id,
-    trust_remote_code=True,
-    torch_dtype=torch.float16,
-    device_map='cuda',
-    # load_in_4bit=True,
-    # bnb_4bit_compute_dtype=torch.float16,
-)
+model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float16).cuda()
 tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
 image_processor = AutoImageProcessor.from_pretrained(model_id)
 
-# ---- setup gradio demo ----
+# ---- gradio demo ----
 model.get_gradio_demo(tokenizer, image_processor).queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
 ```