Update app.py

app.py

@@ -1,33 +1,29 @@
+import os
 import gradio as gr
+from fastai.vision.all import *
+import skimage
 import copy
-import os
-import time
-import llama_cpp
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
-from fastai.vision.all import *
 
-# Load the LLM model
+# Load the FastAI vision model
+learn = load_learner('export.pkl')
+labels = learn.dls.vocab
+
+# Load the Llama language model
 llm = Llama(
     model_path=hf_hub_download(
         repo_id=os.environ.get("REPO_ID", "TheBloke/Llama-2-7B-Chat-GGML"),
         filename=os.environ.get("MODEL_FILE", "llama-2-7b-chat.ggmlv3.q5_0.bin"),
     ),
     n_ctx=2048,
-    n_gpu_layers=50,  # change n_gpu_layers if you have more or less VRAM
+    n_gpu_layers=50,
 )
 
 history = []
-
 system_message = """
-You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe...
+You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  
 """
-# The rest of the system message
-
-# Load the Vision Model
-learn = load_learner('export.pkl')
-
-labels = learn.dls.vocab
 
 def generate_text(message, history):
     temp = ""
@@ -59,36 +55,25 @@ def generate_text(message, history):
         temp += stream["choices"][0]["text"]
         yield temp
 
-    history = ["init", input_prompt]
+    history.append(("USER:", message))
+    history.append(("ASSISTANT:", temp))
 
-def predict(img):
-    try:
-        img = PILImage.create(img)
-    except:
-        return {"bird": "Unknown"}
+# Define the predict function for the FastAI model
+def predict_with_llama_and_generate_text(img):
+    img = PILImage.create(img)
     pred, pred_idx, probs = learn.predict(img)
-    return {"bird": labels[pred_idx], "probs": {labels[i]: float(probs[i]) for i in range(len(labels))}}
-
-title = "Bird Detector with LLM"
-description = "Detect birds and get LLM responses."
-examples = [{"img": "BIRD.png", "message": "Tell me about the bird."}]
-interpretation = 'default'
-enable_queue = True
-
-def combined(img, message):
-    prediction = predict(img)
-    response = list(generate_text(f"I have detected {prediction['bird']} in the image. {message}", history))
-    return response[0]  # Return the first generated response
+    detected_object = labels[pred_idx]
+    
+    response = f"The system has detected {detected_object}. Do you want to know about {detected_object}?"
+    
+    for llama_response in generate_text(response, history):
+        yield llama_response
 
+# Define the Gradio interface
 gr.Interface(
-    fn=combined,
-    inputs=[
-        gr.inputs.Image(),
-        gr.inputs.Textbox(label="Message to LLM")
-    ],
+    fn=predict_with_llama_and_generate_text,
+    inputs=gr.inputs.Image(shape=(512, 512)),
     outputs=gr.outputs.Textbox(),
-    title=title,
-    description=description,
-    examples=examples,
-    interpretation=interpretation,
+    title="Multimodal Assistant",
+    description="An AI model that combines image classification with text generation.",
 ).launch()