Update app.py

app.py

@@ -11,80 +11,85 @@ client = None
 def initialize_client(api_key=None):
     global client
     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(image_path, 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(("Error initializing client", f"{str(e)}"))
+            yield history
+            return
 
-    prompt = "Analyze the attached radiator image and respond with 'accident' if it shows signs of an accident, or 'not accident' if it does 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."
+    prompt = """
+      Determine if the radiator in the image shows signs of being involved in an accident or not
+    """
 
-    messages = [
-        {"role": "system", "content": prompt},
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
-            ],
-        },
-    ]
+    messages = [{"role": "system", "content": prompt}]
+    
+    # Encode the image and add to messages
+    image_base64 = encode_image(image_path)
+    messages.append({
+        "role": "user",
+        "content": [
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/png;base64,{image_base64}"}
+            }
+        ]
+    })
+    history = history + [("Image uploaded", "")]
 
     try:
-        response = client.chat.completions.create(
+        stream = client.chat.completions.create(
             model="meta-llama/Llama-Vision-Free",
             messages=messages,
-            max_tokens=10,
-            temperature=0.0,
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            stream=True,
         )
 
-        # Print the raw response for debugging
-        print("Raw response from model:", response)
+        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
+                history[-1] = ("Image uploaded", response)
+                yield history
 
-        # Extract and clean the response
-        response_text = response.choices[0].message.content.strip().lower()
-        
-        # Return the expected results
-        if "accident" in response_text:
-            return "accident"
-        elif "not accident" in response_text:
-            return "not accident"
-        else:
-            return f"Error: Unexpected response format. Received: {response_text}"
+        if not response:
+            history[-1] = ("Image uploaded", "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. Details: {str(e)}"
+        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] = ("Image uploaded", error_message)
+        yield history
 
-# Gradio interface
+# Set up Gradio interface
 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", label="Upload Radiator Image")
-    result = gr.Textbox(label="Result")
-    submit = gr.Button("Submit")
+    chatbot = gr.Chatbot()
+    img = gr.Image(type="filepath", label="Upload Radiator Image")
+    clear = gr.Button("Clear")
 
-    submit.click(detect_accident, inputs=image_input, outputs=result)
+    img.upload(bot_streaming, inputs=[img, chatbot], outputs=chatbot)
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
     demo.launch(debug=True)