Update app.py (#16)

app.py

@@ -11,23 +11,18 @@ import torch
 import spaces
 import requests
 
-CSS ="""
-.container { display: flex; flex-direction: column; height: 500px; }
-#chatbot { flex-grow: 1; }
-"""
-
 model_id = "xtuner/llava-llama-3-8b-v1_1-transformers"
 
-# processor = AutoProcessor.from_pretrained(model_id)
+processor = AutoProcessor.from_pretrained(model_id)
 
-# model = LlavaForConditionalGeneration.from_pretrained(
-#     model_id, 
-#     torch_dtype=torch.float16, 
-#     low_cpu_mem_usage=True, 
-# )
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True, 
+)
 
-# model.to("cuda:0")
-# model.generation_config.eos_token_id = 128009
+model.to("cuda:0")
+model.generation_config.eos_token_id = 128009
 
 @spaces.GPU
 def bot_streaming(message, history):
@@ -41,34 +36,34 @@ def bot_streaming(message, history):
       if type(hist[0])==tuple:
         image = hist[0][0]
 
-  # if image is None:
-  #     gr.Error("You need to upload an image for LLaVA to work.")
-  # prompt=f"<|start_header_id|>user<|end_header_id|>\n\n<image>\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-  # print(f"prompt: {prompt}")
-  # image = Image.open(image)
-  # inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
+  if image is None:
+      gr.Error("You need to upload an image for LLaVA to work.")
+  prompt=f"<|start_header_id|>user<|end_header_id|>\n\n<image>\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+  print(f"prompt: {prompt}")
+  image = Image.open(image)
+  inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
 
-  # streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True})
-  # generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-  # generated_text = ""
+  streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True})
+  generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+  generated_text = ""
 
-  # thread = Thread(target=model.generate, kwargs=generation_kwargs)
-  # thread.start()
+  thread = Thread(target=model.generate, kwargs=generation_kwargs)
+  thread.start()
 
-  # text_prompt =f"<|start_header_id|>user<|end_header_id|>\n\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-  # print(f"text_prompt: {text_prompt}")
+  text_prompt =f"<|start_header_id|>user<|end_header_id|>\n\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+  print(f"text_prompt: {text_prompt}")
 
-  # buffer = ""
-  # for new_text in streamer:
+  buffer = ""
+  for new_text in streamer:
     
-  #   buffer += new_text
+    buffer += new_text
     
-  #   generated_text_without_prompt = buffer[len(text_prompt):]
-  #   time.sleep(0.04)
-  #   yield generated_text_without_prompt
+    generated_text_without_prompt = buffer[len(text_prompt):]
+    time.sleep(0.04)
+    yield generated_text_without_prompt
 
 
-with gr.Blocks(css=CSS) as demo:
+with gr.Blocks as demo:
     chatbot = gr.ChatInterface(fn=bot_streaming, title="LLaVA Llama-3-8B", examples=[{"text": "What is on the flower?", "files":["./bee.jpg"]},
                                                                       {"text": "How to make this pastry?", "files":["./baklava.png"]}], 
                             description="Try [LLaVA Llama-3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",