Bring back chatbot

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer
+from threading import Thread
 import re
 import time
 from PIL import Image
 import torch
 import spaces
-import subprocess
+#import subprocess
 #subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 
@@ -17,36 +18,37 @@ model = AutoModelForVision2Seq.from_pretrained("HuggingFaceTB/SmolVLM-Instruct",
 
 @spaces.GPU
 def model_inference(
-    images, text, assistant_prefix, decoding_strategy, temperature, max_new_tokens,
+    input_dict, history, decoding_strategy, temperature, max_new_tokens,
     repetition_penalty, top_p
-):
+): 
+    text = input_dict["text"]
+    print(input_dict["files"])
+    if len(input_dict["files"]) > 1:
+      images = [Image.open(image).convert("RGB") for image in input_dict["files"]]
+    elif len(input_dict["files"]) == 1:
+      images = [Image.open(input_dict["files"][0]).convert("RGB")]        
+    
+
     if text == "" and not images:
         gr.Error("Please input a query and optionally image(s).")
 
     if text == "" and images:
         gr.Error("Please input a text query along the image(s).")
 
-    if isinstance(images, Image.Image):
-        images = [images]
+    
 
 
     resulting_messages = [
                 {
                     "role": "user",
-                    "content": [{"type": "image"}] + [
+                    "content": [{"type": "image"} for _ in range(len(images))] + [
                         {"type": "text", "text": text}
                     ]
                 }
             ]
-
-    if assistant_prefix:
-      text = f"{assistant_prefix} {text}"
-
-
     prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
     inputs = processor(text=prompt, images=[images], return_tensors="pt")
     inputs = {k: v.to("cuda") for k, v in inputs.items()}
-
     generation_args = {
         "max_new_tokens": max_new_tokens,
         "repetition_penalty": repetition_penalty,
@@ -65,119 +67,82 @@ def model_inference(
         generation_args["top_p"] = top_p
 
     generation_args.update(inputs)
-
     # Generate
-    generated_ids = model.generate(**generation_args)
-
-    generated_texts = processor.batch_decode(generated_ids[:, generation_args["input_ids"].size(1):], skip_special_tokens=True)
-    return generated_texts[0]
-
-
-with gr.Blocks(fill_height=False) as demo:
-    gr.Markdown("## SmolVLM: Small yet Mighty 💫")
-    gr.Markdown("Play with [HuggingFaceTB/SmolVLM-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM-Instruct) in this demo. To get started, upload an image and text or try one of the examples.")
-    with gr.Column():
-        with gr.Row():
-            image_input = gr.Image(label="Upload your Image", type="pil")
-
-            with gr.Column():
-                query_input = gr.Textbox(label="Prompt")
-                assistant_prefix = gr.Textbox(label="Assistant Prefix", placeholder="Let's think step by step.")
-
-                submit_btn = gr.Button("Submit")
-        output = gr.Textbox(label="Output")
-
-        
-        with gr.Accordion(label="Advanced Generation Parameters", open=False):
-            examples=[
-                    ["example_images/rococo.jpg", "What art era is this?", "", "Top P Sampling", 0.4, 512, 1.2, 0.8],
-                    ["example_images/examples_wat_arun.jpg", "I'm planning a visit to this temple, give me travel tips.",  "", "Greedy", 0.4, 512, 1.2, 0.8],
-                    ["example_images/examples_invoice.png", "What is the due date and the invoice date?",  "", "Top P Sampling", 0.4, 512, 1.2, 0.8],
-                    ["example_images/s2w_example.png", "What is this UI about?",  "", "Top P Sampling", 0.4, 512, 1.2, 0.8],
-                    ["example_images/examples_weather_events.png", "Where do the severe droughts happen according to this diagram?",  "", "Top P Sampling", 0.4, 512, 1.2, 0.8],
-            ]
-            # Hyper-parameters for generation
-            max_new_tokens = gr.Slider(
-                minimum=8,
-                maximum=1024,
-                value=512,
-                step=1,
-                interactive=True,
-                label="Maximum number of new tokens to generate",
-            )
-            repetition_penalty = gr.Slider(
-                minimum=0.01,
-                maximum=5.0,
-                value=1.2,
-                step=0.01,
-                interactive=True,
-                label="Repetition penalty",
-                info="1.0 is equivalent to no penalty",
-            )
-            temperature = gr.Slider(
-                minimum=0.0,
-                maximum=5.0,
-                value=0.4,
-                step=0.1,
-                interactive=True,
-                label="Sampling temperature",
-                info="Higher values will produce more diverse outputs.",
-            )
-            top_p = gr.Slider(
-                minimum=0.01,
-                maximum=0.99,
-                value=0.8,
-                step=0.01,
-                interactive=True,
-                label="Top P",
-                info="Higher values is equivalent to sampling more low-probability tokens.",
-            )
-            decoding_strategy = gr.Radio(
-                [
-                    "Top P Sampling",
-                    "Greedy",
-                    
-                ],
-                value="Top P Sampling",
-                label="Decoding strategy",
-                interactive=True,
-                info="Higher values is equivalent to sampling more low-probability tokens.",
-            )
-            decoding_strategy.change(
-                fn=lambda selection: gr.Slider(
-                    visible=(
-                        selection in ["contrastive_sampling", "beam_sampling", "Top P Sampling", "sampling_top_k"]
-                    )
-                ),
-                inputs=decoding_strategy,
-                outputs=temperature,
-            )
-
-            decoding_strategy.change(
-                fn=lambda selection: gr.Slider(
-                    visible=(
-                        selection in ["contrastive_sampling", "beam_sampling", "Top P Sampling", "sampling_top_k"]
-                    )
-                ),
-                inputs=decoding_strategy,
-                outputs=repetition_penalty,
-            )
-            decoding_strategy.change(
-                fn=lambda selection: gr.Slider(visible=(selection in ["Top P Sampling"])),
-                inputs=decoding_strategy,
-                outputs=top_p,
-            )
-        gr.Examples(
-                        examples = examples,
-                        inputs=[image_input, query_input, assistant_prefix, decoding_strategy, temperature,
-                                                              max_new_tokens, repetition_penalty, top_p],
-                        outputs=output,
-                        fn=model_inference
-                    )   
-        
-
-        submit_btn.click(model_inference, inputs = [image_input, query_input, assistant_prefix, decoding_strategy, temperature,
-                                                      max_new_tokens, repetition_penalty, top_p], outputs=output)
-
-
-demo.launch(debug=True)
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens= True)
+    generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
+
+    thread = Thread(target=model.generate, kwargs=generation_args)
+    thread.start()
+    thread.join()
+    
+    buffer = ""
+    
+      
+    for new_text in streamer:
+    
+      buffer += new_text
+      generated_text_without_prompt = buffer#[len(ext_buffer):]
+      time.sleep(0.01)
+      yield buffer
+
+
+examples=[
+              [{"text": "What art era do these artpieces belong to?", "files": ["example_images/rococo.jpg", "example_images/rococo_1.jpg"]}, "Greedy", 0.4, 512, 1.2, 0.8],
+              [{"text": "I'm planning a visit to this temple, give me travel tips.", "files": ["example_images/examples_wat_arun.jpg"]},  "Greedy", 0.4, 512, 1.2, 0.8],
+              [{"text":  "What is the due date and the invoice date?", "files": ["example_images/examples_invoice.png"]}, "Greedy", 0.4, 512, 1.2, 0.8],
+              [{"text": "What is this UI about?", "files": ["example_images/s2w_example.png"]},   "Greedy", 0.4, 512, 1.2, 0.8],
+              [{"text":  "Where do the severe droughts happen according to this diagram?", "files": ["example_images/examples_weather_events.png"]},  "Greedy", 0.4, 512, 1.2, 0.8],
+      ]
+demo = gr.ChatInterface(fn=model_inference, title="SmolVLM: Small yet Mighty 💫", 
+                description="Play with [HuggingFaceTB/SmolVLM-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM-Instruct) in this demo. To get started, upload an image and text or try one of the examples. This checkpoint works best with single turn conversations, so clear the conversation after a single turn.",
+                examples=examples,
+                textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"), stop_btn="Stop Generation", multimodal=True, 
+                        additional_inputs=[gr.Radio(["Top P Sampling",
+              "Greedy"],
+          value="Greedy",
+          label="Decoding strategy",
+          #interactive=True,
+          info="Higher values is equivalent to sampling more low-probability tokens.",
+                                                  
+      ), gr.Slider(
+          minimum=0.0,
+          maximum=5.0,
+          value=0.4,
+          step=0.1,
+          interactive=True,
+          label="Sampling temperature",
+          info="Higher values will produce more diverse outputs.",
+      ),
+                                            gr.Slider(
+          minimum=8,
+          maximum=1024,
+          value=512,
+          step=1,
+          interactive=True,
+          label="Maximum number of new tokens to generate",
+      ), gr.Slider(
+          minimum=0.01,
+          maximum=5.0,
+          value=1.2,
+          step=0.01,
+          interactive=True,
+          label="Repetition penalty",
+          info="1.0 is equivalent to no penalty",
+      ), 
+         gr.Slider(
+          minimum=0.01,
+          maximum=0.99,
+          value=0.8,
+          step=0.01,
+          interactive=True,
+          label="Top P",
+          info="Higher values is equivalent to sampling more low-probability tokens.",
+      )],cache_examples=False
+                )
+     
+      
+      
+
+demo.launch(debug=True)
+