Update app.py

app.py

@@ -9,11 +9,9 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer
+from transformers import AutoModel, AutoProcessor, TextIteratorStreamer
 import os
 import time
-from huggingface_hub import hf_hub_download
-
 
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
@@ -24,7 +22,7 @@ MODEL_NAME = MODEL_ID.split("/")[-1]
 
 TITLE = "<h1><center>VL-Chatbox</center></h1>"
 
-DESCRIPTION = "<h3><center>MODEL: " + MODEL_NAME + "</center></h3>"
+DESCRIPTION = "<h3><center>MODEL: " + f'[{MODEL_NAME}](https://hf.co/models/{MODEL_NAME})' + "</center></h3>"
 
 CSS = """
 .duplicate-button {
@@ -35,15 +33,13 @@ CSS = """
 }
 """
 
-model = AutoModelForCausalLM.from_pretrained(
+model = AutoModel.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
     trust_remote_code=True
 ).to(0)
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-eos_token_id=processor.tokenizer.eos_token_id
-
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model.eval()
 
 
 
@@ -53,8 +49,8 @@ def stream_chat(message, history: list, temperature: float, max_new_tokens: int)
     print(f'history is - {history}')
     conversation = [] 
     if message["files"]:
-        image = Image.open(message["files"][-1])
-        conversation.append({"role": "user", "content": f"<|image_1|>\n{message['text']}"})
+        image = Image.open(message["files"][-1]).convert('RGB')
+        conversation.append({"role": "user", "content": message['text']})
     else:
         if len(history) == 0:
             raise gr.Error("Please upload an image first.")
@@ -62,29 +58,29 @@ def stream_chat(message, history: list, temperature: float, max_new_tokens: int)
         else:
             image = Image.open(history[0][0][0])
             for prompt, answer in history:
-                if answer is None:
-                    conversation.extend([{"role": "user", "content":"<|image_1|>"},{"role": "assistant", "content": ""}])
-                else:
+ #               if answer is None:
+ #                   conversation.extend([{"role": "user", "content":"<|image_1|>"},{"role": "assistant", "content": ""}])
+ #               else:
                     conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
             conversation.append({"role": "user", "content": message['text']})    
     print(f"Conversation is -\n{conversation}")
-    inputs = processor.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
-    inputs_ids = processor(inputs, image, return_tensors="pt").to(0)
-    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,}) 
+
+    streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,}) 
 
     generate_kwargs = dict(
+        image=image,
+        msg=conversation,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
-        do_sample=True,
-        eos_token_id=eos_token_id,
+        sampling=True,
+        tokenizer=tokenizer,
     )
     if temperature == 0:
-        generate_kwargs["do_sample"] = False
-    generate_kwargs = {**inputs_ids, **generate_kwargs}
+        generate_kwargs["sampling"] = False
     
 
-    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread = Thread(target=model.chat, kwargs=generate_kwargs)
     thread.start()
 
     buffer = ""