Update app.py

app.py

@@ -28,70 +28,56 @@ processor = AutoProcessor.from_pretrained(model_id)
 # Confirming and setting the eos_token_id (if necessary)
 model.generation_config.eos_token_id = processor.tokenizer.eos_token_id
 
-@spaces.GPU(duration=120)
-def krypton(input, history):
-    print(f"Input: {input}")  # Debug input
-    print(f"History: {history}")  # Debug history
-
-    image_path = None
-    if input["files"]:
-        print("Found the image")
-        image_path = input["files"][-1]["path"] if isinstance(input["files"][-1], dict) else input["files"][-1]
-        print(f"Image path: {image_path}")
+@spaces.GPU
+def bot_streaming(message, history):
+    print(message)
+    if message["files"]:
+        # message["files"][-1] is a Dict or just a string
+        if type(message["files"][-1]) == dict:
+            image = message["files"][-1]["path"]
+        else:
+            image = message["files"][-1]
     else:
+        # if there's no image uploaded for this turn, look for images in the past turns
+        # kept inside tuples, take the last one
         for hist in history:
-            if isinstance(hist[0], tuple):
-                image_path = hist[0][0]
-                break
+            if type(hist[0]) == tuple:
+                image = hist[0][0]
+    try:
+        if image is None:
+            # Handle the case where image is None
+            gr.Error("You need to upload an image for LLaVA to work.")
+    except NameError:
+        # Handle the case where 'image' is not defined at all
+        gr.Error("You need to upload an image for LLaVA to work.")
 
-    if not image_path:
-        gr.Error("You need to upload an image for Krypton to work.")
-        return
+    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)
 
-    try:
-        image = Image.open(image_path)
-        image.show()  # Show the image to confirm it's loaded
-        print(f"Image open: {image}")
-    except Exception as e:
-        print(f"Error opening image: {e}")
-        gr.Error("Failed to open the image.")
-        return
-
-    # Adding more context to the prompt with a placeholder for the image
-    prompt = f"user: Here is an image and a question about it.\n<image>{input['text']}\nassistant: "
-    print("Made the prompt")
+    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": False, "skip_prompt": True})
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False)
 
-    try:
-        inputs = processor(text=prompt, images=image, return_tensors='pt').to('cuda', torch.float16)
-        print(f"Processed inputs: {inputs}")
-    except Exception as e:
-        print(f"Error processing inputs: {e}")
-        gr.Error("Failed to process the inputs.")
-        return
-    
-    # Streamer
-    print('About to init streamer')
-    streamer = TextIteratorStreamer(processor.tokenizer, skip_special_tokens=False, skip_prompt=True)
-    
-    # Generation kwargs
-    generation_kwargs = dict(
-        inputs=inputs['input_ids'],
-        attention_mask=inputs['attention_mask'],
-        streamer=streamer,
-        max_new_tokens=1024,
-        do_sample=False
-    )
-    
-    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-    print('Thread about to 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}")
+
     buffer = ""
-    # time.sleep(0.5)
+    time.sleep(0.5)
     for new_text in streamer:
+        # find <|eot_id|> and remove it from the new_text
+        if "<|eot_id|>" in new_text:
+            new_text = new_text.split("<|eot_id|>")[0]
         buffer += new_text
+
+        # generated_text_without_prompt = buffer[len(text_prompt):]
         generated_text_without_prompt = buffer
-        # time.sleep(0.06)
+        # print(generated_text_without_prompt)
+        time.sleep(0.06)
+        # print(f"new_text: {generated_text_without_prompt}")
         yield generated_text_without_prompt