Update app.py

app.py

@@ -10,70 +10,143 @@ client = None
 
 def initialize_client(api_key=None):
     global client
+    # Fetch the API key from the environment if it's not provided directly
     api_key = os.getenv("TOGETHER_API_KEY")
-    if api_key:
-        client = Together()
-        print("Client initialized successfully.")
-    else:
-        print("API key not found. Please set TOGETHER_API_KEY environment variable.")
+    print(api_key)
+    
+    client = Together()
+
 
 def encode_image(image_path):
-    try:
-        with Image.open(image_path) as img:
-            buffered = io.BytesIO()
-            img.save(buffered, format="PNG")
-            return base64.b64encode(buffered.getvalue()).decode("utf-8")
-    except Exception as e:
-        print(f"Error encoding image: {str(e)}")
-        return None
+    with Image.open(image_path) as img:
+        buffered = io.BytesIO()
+        img.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-def detect_accident(image_path):
+
+def bot_streaming(message, history):
+    max_new_tokens=250 
+    temperature=0.7
     if client is None:
-        initialize_client()
-        if client is None:
-            return "Error: Client not initialized."
+        try:
+            initialize_client()
+        except Exception as e:
+            history.append((message, f"Error initializing client: {str(e)}"))
+            yield history
+            return
 
-    prompt = "Determine if the radiator in the image shows signs of being involved in an accident or not."
-    image_base64 = encode_image(image_path)
-    
-    if not image_base64:
-        return "Error: Could not encode the image. Please try again with a different one."
-
-    messages = [
-        {"role": "system", "content": prompt},
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
-            ],
-        },
-    ]
+    prompt = """
+      Determine if the radiator in the image shows signs of being involved in an accident or not
+"""
+
+
+    messages = [{"role": "system", "content": prompt}]
+
+    # Add history to messages
+    for user_msg, assistant_msg in history:
+        if isinstance(user_msg, str):  # Text message
+            messages.append(
+                {"role": "user", "content": [{"type": "text", "text": user_msg}]}
+            )
+        elif isinstance(user_msg, dict):  # Image message
+            image_base64 = encode_image(user_msg["image_path"])
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": user_msg.get("text", "")},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{image_base64}"
+                            },
+                        },
+                    ],
+                }
+            )
+        messages.append(
+            {"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]}
+        )
+
+    # Prepare the current message
+    user_message_content = []
+    if isinstance(message, dict):
+        if message.get("text"):
+            user_message_content.append({"type": "text", "text": message["text"]})
+        if message.get("files") and len(message["files"]) > 0:
+            image_path = message["files"][0]
+            image_base64 = encode_image(image_path)
+            user_message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{image_base64}"},
+                }
+            )
+    elif isinstance(message, str):
+        user_message_content.append({"type": "text", "text": message})
+
+    current_message = {"role": "user", "content": user_message_content}
+    messages.append(current_message)
+
+    # Add the user's message to the history
+    user_display_message = message["text"] if isinstance(message, dict) else message
+    history = history + [(user_display_message, "")]
 
     try:
-        response = client.chat.completions.create(
+        stream = client.chat.completions.create(
             model="meta-llama/Llama-Vision-Free",
             messages=messages,
-            max_tokens=10,
-            temperature=0.5
-        ).choices[0].message["content"].strip()
-        
-        # Check if response contains "accident" or "not accident"
-        return "accident" if "accident" in response.lower() else "not accident"
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            stream=True,
+        )
+
+        response = ""
+        for chunk in stream:
+            if (
+                chunk.choices
+                and chunk.choices[0].delta
+                and chunk.choices[0].delta.content is not None
+            ):
+                response += chunk.choices[0].delta.content
+                # Update the assistant's response in the history
+                history[-1] = (user_display_message, response)
+                yield history
+
+        if not response:
+            history[-1] = (
+                user_display_message,
+                "No response generated. Please try again.",
+            )
+            yield history
 
     except Exception as e:
-        print(f"Error during model call: {str(e)}")
-        return f"Error: Could not process the image. Please try again with a different one."
+        error_message = (
+            "The image is too large. Please try with a smaller image or compress the existing one."
+            if "Request Entity Too Large" in str(e)
+            else f"An error occurred: {str(e)}"
+        )
+        history[-1] = (user_display_message, error_message)
+        yield history
+
 
-# Gradio interface
+# The rest of your Gradio interface code remains the same
 with gr.Blocks() as demo:
     gr.Markdown("# Radiator Accident Detection")
-    gr.Markdown("Upload an image of a radiator to determine if it shows signs of an accident.")
+    gr.Markdown(
+        "Upload an image of a radiator to determine if it shows signs of an accident"
+    )
 
-    image_input = gr.Image(type="filepath")
-    result = gr.Textbox(label="Result")
-    submit = gr.Button("Submit")
+    chatbot = gr.Chatbot()
+    msg = gr.MultimodalTextbox(label="")
+    clear = gr.Button("Clear")
 
-    submit.click(detect_accident, inputs=image_input, outputs=result)
+    msg.submit(
+        bot_streaming,
+        [msg, chatbot],
+        chatbot,
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)