Update app.py

app.py

@@ -1,10 +1,4 @@
 import gradio as gr
-import subprocess
-# subprocess.run(
-#     "pip install flash-attn --no-build-isolation",
-#     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-#     shell=True,
-# )
 from llava.model.builder import load_pretrained_model
 from llava.mm_utils import process_images, tokenizer_image_token
 from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
@@ -13,7 +7,6 @@ from PIL import Image
 import copy
 import torch
 import warnings
-import requests
 
 warnings.filterwarnings("ignore")
 
@@ -26,41 +19,26 @@ device_map = "auto"
 tokenizer, model, image_processor, max_length = load_pretrained_model(pretrained, None, model_name, device_map=device_map, attn_implementation="sdpa")
 model.eval()
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    image=None,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    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 image:
+def respond(image_path, question, temperature, max_tokens):
+    try:
         # Load and process the image
-        if isinstance(image, str):
-            image = Image.open(image)
-
+        image = Image.open(image_path)
         image_tensor = process_images([image], image_processor, model.config)
         image_tensor = [_image.to(dtype=torch.float16, device=device) for _image in image_tensor]
 
+        # Prepare the conversation template
         conv_template = "qwen_1_5"
-        question = DEFAULT_IMAGE_TOKEN + "\n" + message
+        formatted_question = DEFAULT_IMAGE_TOKEN + "\n" + question
         conv = copy.deepcopy(conv_templates[conv_template])
-        conv.append_message(conv.roles[0], question)
+        conv.append_message(conv.roles[0], formatted_question)
         conv.append_message(conv.roles[1], None)
         prompt_question = conv.get_prompt()
 
+        # Tokenize input
         input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
         image_sizes = [image.size]
 
+        # Generate response
         cont = model.generate(
             input_ids,
             images=image_tensor,
@@ -70,45 +48,28 @@ def respond(
             max_new_tokens=max_tokens,
         )
 
-        response = tokenizer.batch_decode(cont, skip_special_tokens=True)[0]
-    else:
-        messages.append({"role": "user", "content": message})
-
-        conv_template = "qwen_1_5"
-        conv = copy.deepcopy(conv_templates[conv_template])
-        conv.append_message(conv.roles[0], message)
-        conv.append_message(conv.roles[1], None)
-        prompt_question = conv.get_prompt()
-
-        input_ids = tokenizer(prompt_question, return_tensors="pt", max_length=max_tokens, truncation=True).to(device)
-
-        cont = model.generate(
-            input_ids,
-            do_sample=True,
-            temperature=temperature,
-            max_new_tokens=max_tokens,
-            top_p=top_p,
-        )
-
-        response = tokenizer.batch_decode(cont, skip_special_tokens=True)[0]
-
-    yield response
-
-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"),
+        text_outputs = tokenizer.batch_decode(cont, skip_special_tokens=True)
+        return text_outputs[0]
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Gradio Interface
+def chat_interface(image, question, temperature, max_tokens):
+    if not image or not question:
+        return "Please provide both an image and a question."
+    return respond(image.name, question, temperature, max_tokens)
+
+demo = gr.Interface(
+    fn=chat_interface,
+    inputs=[
+        gr.Image(type="file", label="Input Image"),
+        gr.Textbox(label="Question"),
         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)",
-        ),
-        gr.Image(type="filepath", label="Input Image (optional)"),
+        gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max Tokens"),
     ],
+    outputs="text",
+    title="AI-Safeguard Ivy-VL-Llava Image Question Answering",
+    description="Upload an image and ask a question about it. The model will provide a response based on the visual and textual input."
 )
 
 if __name__ == "__main__":