Update app.py

app.py

@@ -1,64 +1,139 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer
+from threading import Thread
+import re
+import time
+from PIL import Image
+import torch
+import spaces
+#import subprocess
+#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-Instruct")
+model = AutoModelForVision2Seq.from_pretrained("HuggingFaceTB/SmolVLM-Instruct", 
+        torch_dtype=torch.bfloat16,
+        #_attn_implementation="flash_attention_2"
+        ).to("cuda")
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+@spaces.GPU
+def model_inference(
+    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")]        
+    
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    if text == "" and not images:
+        gr.Error("Please input a query and optionally image(s).")
 
-    messages.append({"role": "user", "content": message})
+    if text == "" and images:
+        gr.Error("Please input a text query along the image(s).")
 
-    response = ""
+    
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
 
-        response += token
-        yield response
+    resulting_messages = [
+                {
+                    "role": "user",
+                    "content": [{"type": "image"} for _ in range(len(images))] + [
+                        {"type": "text", "text": 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,
 
+    }
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+    assert decoding_strategy in [
+        "Greedy",
+        "Top P Sampling",
+    ]
+    if decoding_strategy == "Greedy":
+        generation_args["do_sample"] = False
+    elif decoding_strategy == "Top P Sampling":
+        generation_args["temperature"] = temperature
+        generation_args["do_sample"] = True
+        generation_args["top_p"] = top_p
 
+    generation_args.update(inputs)
+    # Generate
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens= True)
+    generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
 
-if __name__ == "__main__":
-    demo.launch()
+    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
+
+
+demo = gr.ChatInterface(fn=model_inference, title="Geoscience AI Interpreter", 
+                description="This app take the thin sections, seismic images etc. and interpret them. You just upload an image and text along with it. It works best with single turn conversations, so clear the conversation after a single turn.",
+                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)