Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoProcessor,TextIteratorStreamer
 import os
 import time
 from huggingface_hub import hf_hub_download
@@ -18,7 +18,7 @@ MODEL_NAME = MODEL_ID.split("/")[-1]
 
 TITLE = "<h1><center>VL-Chatbox</center></h1>"
 
-DESCRIPTION = "<h3><center>MODEL LOADED: " + MODEL_NAME + "</center></h3>"
+DESCRIPTION = "<h3><center>MODEL: " + MODEL_NAME + "</center></h3>"
 
 DEFAULT_SYSTEM = "You named Chatbox. You are a good assitant."
 
@@ -31,54 +31,15 @@ CSS = """
 }
 """
 
-filenames = [
-    "config.json",
-    "generation_config.json",
-    "model-00001-of-00004.safetensors",
-    "model-00002-of-00004.safetensors",
-    "model-00003-of-00004.safetensors",
-    "model-00004-of-00004.safetensors",
-    "model.safetensors.index.json",
-    "special_tokens_map.json",
-    "tokenizer.json",
-    "tokenizer_config.json"
-]
-
-for filename in filenames:
-    downloaded_model_path = hf_hub_download(
-        repo_id=MODEL_ID,
-        filename=filename,
-        local_dir="./model/"
-    )
-
-for items in os.listdir("./model"):
-    print(items)
-
-# def no_logger():
-#     logging.config.dictConfig({
-#         'version': 1,
-#         'disable_existing_loggers': True,
-#     })
-
-
 model = AutoModelForCausalLM.from_pretrained(
-    "./model/",
+    MODEL_ID,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     trust_remote_code=True
 ).to(0)
-tokenizer = AutoTokenizer.from_pretrained("./model/",trust_remote_code=True)
-vision_tower = model.get_vision_tower()
-vision_tower.load_model()
-vision_tower.to(device="cuda", dtype=torch.float16)
-image_processor = vision_tower.image_processor
-tokenizer.pad_token = tokenizer.eos_token
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# Define terminators (if applicable, adjust as needed)
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
+eos_token_id=processor.tokenizer.eos_token_id
 
 
 
@@ -88,49 +49,36 @@ def stream_chat(message, history: list, system: str, temperature: float, max_new
     print(message)
     conversation = [{"role": "system", "content": system or DEFAULT_SYSTEM}]
     for prompt, answer in history:
-        conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
-
+        conversation.extend([{"role": "user", "content": f"<|image_1|>\n{prompt}"}, {"role": "assistant", "content": answer}])
     conversation.append({"role": "user", "content": message['text']})
+    
     if message["files"]:
         image = Image.open(message["files"][0]).convert('RGB')
-        # Process the conversation text
-        inputs = model.build_conversation_input_ids(
-            tokenizer, 
-            query=message['text'], 
-            image=image, 
-            image_processor=image_processor,
-        )
-        input_ids = inputs["input_ids"].to(device='cuda', non_blocking=True)
-        images = inputs["image"].to(dtype=torch.float16, device='cuda', non_blocking=True)
     else:
-        input_ids = tokenizer.apply_chat_template(
-            conversation, 
-            add_generation_prompt=True, 
-            return_tensors="pt"
-        ).to(model.device)
-        images = None
+        image = None
+        
+    prompt = processor.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    inputs = processor(prompt, [image], return_tensors="pt").to(0) 
 
     generate_kwargs = dict(
-        input_ids=input_ids,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         do_sample=True,
-        num_beams=1,
-        eos_token_id=terminators,
-        images=images
+        eos_token_id=eos_token_id,
     )
     if temperature == 0:
         generate_kwargs["do_sample"] = False
+    generate_kwargs = {**inputs, **generate_kwargs}
+    
+    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,}) 
 
-    output_ids=model.generate(**generate_kwargs)
-    input_token_len = input_ids.shape[1]
-    outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
-    outputs = outputs.strip()
-
-    for i in range(len(outputs)):
-        time.sleep(0.05)
-        yield outputs[: i + 1]
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
 
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
 
 
 chatbot = gr.Chatbot(height=450)