placed processor tokenizer

app.py

@@ -19,107 +19,50 @@ model = LlavaForConditionalGeneration.from_pretrained(
     model_id,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True
-).to('cuda:0')
+)
 
 processor = AutoProcessor.from_pretrained(model_id)
 
-model.generation_config.eos_token_id = 128009
+# 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):
-    """
-    Recieves inputs (prompts with images if they were added),
-    the image is formated for pil and prompt is formated for the model,
-    to place it's output to the user, these prompts and images are passed in
-    the processor and generation of the model, than the output is decoded from the processor,
-    onto the UI.
-    """
+def krypton(input, history):
     if input["files"]:
-        if type(input["files"][-1]) == dict:
-            image = input["files"][-1]["path"]
-        else:
-            image = input["files"][-1]
+        image = input["files"][-1]["path"] if isinstance(input["files"][-1], dict) else input["files"][-1]
     else:
-        # If no images were passed now, look at the past images to keep up as reference still to the prompts
-        # kept inside in tuples, the last one
+        image = None
         for hist in history:
-            if type(hist[0]) == tuple:
+            if isinstance(hist[0], tuple):
                 image = hist[0][0]
-    try:
-        if image is None:
-            gr.Error("You need to upload an image please for krypton to work.")
-    except NameError:
-        # Image is not defined at all
-        gr.Error("Uplaod an image for Krypton to work")
-
-    prompt = ("<|start_header_id|>user<|end_header_id|>\n\n<image>\n{input['text']}<|eot_id|>"
-              "<|start_header_id|>assistant<|end_header_id|>\n\n")
     
+    if not image:
+        gr.Error("You need to upload an image for Krypton to work.")
+        return
+    
+    prompt = f"user\n\n<image>\n{input['text']}\nassistant\n\n"
     image = Image.open(image)
-    inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
-
+    inputs = processor(prompt, images=image, return_tensors='pt').to(0, torch.float16)
+    
     # Streamer
-    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": False, "skip_prompt": True})
-
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_special_tokens=False, skip_prompt=True)
+    
     # Generation kwargs
     generation_kwargs = dict(
-        inputs=inputs,
+        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)
     thread.start()
-
+    
     buffer = ""
     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
-        # print(generated_text_without_prompt)
         time.sleep(0.06)
-        # print(f"new_text: {generated_text_without_prompt}")
-        yield generated_text_without_prompt
-
-chatbot=gr.Chatbot(height=600, label="Krypt AI")
-chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter your question or upload an image.", show_label=False)
-with gr.Blocks(fill_height=True) as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.ChatInterface(
-        fn=krypton,
-        chatbot=chatbot,
-        fill_height=True,
-        # additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
-        # additional_inputs=[
-        #     gr.Slider(minimum=20,
-        #               maximum=80,
-        #               step=1,
-        #               value=50,
-        #               label="Max New Tokens",
-        #               render=False),
-        #     gr.Slider(minimum=0.0,
-        #               maximum=1.0,
-        #               step=0.1,
-        #               value=0.7,
-        #               label="Temperature",
-        #               render=False),
-        #     gr.Slider(minimum=1,
-        #               maximum=12,
-        #               step=1,
-        #               value=5,
-        #               label="Number of Beams",
-        #               render=False),
-        # ],
-        multimodal=True,
-        textbox=chat_input,
-    )
-
-demo.queue(api_open=False)
-demo.launch(show_api=False, share=False)
+        yield generated_text_without_prompt