Update app.py

app.py

@@ -2,13 +2,12 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
 
 
-os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-MODEL_LIST = ["openbmb/MiniCPM-Llama3-V-2_5","openbmb/MiniCPM-Llama3-V-2_5-int4"]
+MODEL_LIST = ["THUDM/glm-4v-9b"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = os.environ.get("MODEL_ID")
 MODEL_NAME = MODEL_ID.split("/")[-1]
@@ -28,7 +27,8 @@ CSS = """
 
 model = AutoModel.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
     trust_remote_code=True
 ).to(0)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
@@ -36,45 +36,48 @@ model.eval()
 
 
 @spaces.GPU()
-def stream_chat(message, history: list, temperature: float, max_new_tokens: int):
+def stream_chat(message, history: list, temperature: float, max_length: int):
     print(f'message is - {message}')
     print(f'history is - {history}')
     conversation = []
     if message["files"]:
         image = Image.open(message["files"][-1]).convert('RGB')
-        conversation.append({"role": "user", "content": message['text']})
+        conversation.append({"role": "user", "image": image, "content": message['text']})
     else:
         if len(history) == 0:
             raise gr.Error("Please upload an image first.")
             image = None
         else:
-            image = Image.open(history[0][0][0])
+            image = Image.open(history[0][1])
             for prompt, answer in history:
                 if answer is None:
                     conversation.extend([{"role": "user", "content": prompt},{"role": "assistant", "content": ""}])
                 else:
                     conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
-            conversation.append({"role": "user", "content": message['text']})
+            conversation.append({"role": "user", "image": image, "content": message['text']})
     print(f"Conversation is -\n{conversation}")
 
+    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+
     generate_kwargs = dict(
-        image=image,
-        msgs=conversation,
-        max_new_tokens=max_new_tokens,
+        max_length=max_length,
+        streamer=streamer,
+        do_sample=True,
+        top_k=1,
         temperature=temperature,
-        sampling=True,
-        tokenizer=tokenizer,
-        stream=True
+        repetition_penalty=1.2,
     )
-    if temperature == 0:
-        generate_kwargs["sampling"] = False
-
-    response = model.chat(**generate_kwargs)
+    gen_kwargs = {**input_ids, **generate_kwargs}
 
-    buffer = ""
-    for new_text in response:
-        buffer += new_text
-        yield buffer
+    with torch.no_grad():
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            yield buffer
+ 
 
 
 
@@ -117,7 +120,7 @@ with gr.Blocks(css=CSS) as demo:
                 maximum=4096,
                 step=1,
                 value=1024,
-                label="Max new tokens",
+                label="Max Length",
                 render=False,
             ),
         ],