Update app.py

app.py

@@ -1,122 +1,70 @@
-import os
 import gradio as gr
+import os
 from together import Together
-from PIL import Image
-import io
-import base64
+from typing import List, Tuple
 
-# Initialize the Together AI client
-client = Together(api_key=os.environ.get('TOGETHER_API_KEY'))
+# Initialize Together client
+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: {e}")
-        return None
+# Ensure API key is set
+if "TOGETHER_API_KEY" not in os.environ:
+    raise ValueError("Please set the TOGETHER_API_KEY environment variable")
 
-def chat_with_image(message, image, history):
-    # Prepare the messages
-    messages = [{"role": "system", "content": "You are a helpful assistant that can analyze images and text."}]
-    
-    for human, assistant in history:
-        if human.startswith("Image: "):
-            # This is an image message
-            image_path = human.split(": ", 1)[1]
-            encoded_image = encode_image(image_path)
-            if encoded_image:
-                messages.append({
-                    "role": "user",
-                    "content": f"[IMAGE]{encoded_image}[/IMAGE]\nWhat is in this image?"
-                })
-            else:
-                messages.append({"role": "user", "content": "I tried to upload an image, but there was an error."})
-        else:
-            # This is a text-only message
-            messages.append({"role": "user", "content": human})
-        messages.append({"role": "assistant", "content": assistant})
-    
-    # Add the current message
-    if image:
-        encoded_image = encode_image(image)
-        if encoded_image:
-            messages.append({
-                "role": "user",
-                "content": f"[IMAGE]{encoded_image}[/IMAGE]\n{message or 'What is in this image?'}"
-            })
-        else:
-            messages.append({"role": "user", "content": "I tried to upload an image, but there was an error."})
-    else:
-        messages.append({"role": "user", "content": message})
-    
-    # Call the Together AI API
-    try:
-        response = client.chat.completions.create(
-            model="meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo",
-            messages=messages,
-            max_tokens=512,
-            temperature=0.7,
-            top_p=0.7,
-            top_k=50,
-            repetition_penalty=1,
-            stop=["<|eot_id|>", "<|eom_id|>"],
-            stream=True
-        )
-        
-        # Accumulate the response
-        full_response = ""
-        for chunk in response:
-            if chunk.choices[0].delta.content is not None:
-                full_response += chunk.choices[0].delta.content
-                yield full_response
-    except Exception as e:
-        # Enhanced error handling
-        import traceback
-        traceback.print_exc()
-        if hasattr(e, 'response') and e.response is not None:
-            try:
-                error_content = e.response.json()
-                print("Error response JSON:", error_content)
-            except Exception:
-                print("Error response text:", e.response.text)
-        yield f"An error occurred: {str(e)}"
+def call_llama_vision_api(prompt: str, image_url: str) -> str:
+    getDescriptionPrompt = "You are a UX/UI designer. Describe the attached screenshot or UI mockup in detail. I will feed in the output you give me to a coding model that will attempt to recreate this mockup, so please think step by step and describe the UI in detail. Pay close attention to background color, text color, font size, font family, padding, margin, border, etc. Match the colors and sizes exactly. Make sure to mention every part of the screenshot including any headers, footers, etc. Use the exact text from the screenshot."
 
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": getDescriptionPrompt + "\n\n" + prompt},
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": image_url,
+                    },
+                },
+            ],
+        }
+    ]
 
+    stream = client.chat.completions.create(
+        model="meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo",
+        messages=messages,
+        stream=True,
+    )
 
+    response = ""
+    for chunk in stream:
+        content = chunk.choices[0].delta.content or ""
+        response += content
+        yield response
 
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    image = gr.Image(type="filepath")
-    clear = gr.Button("Clear")
+def chat(message: str, history: List[Tuple[str, str]], image_url: str) -> Tuple[str, List[Tuple[str, str]]]:
+    if not message:
+        return "", history
 
-    def user(user_message, image, history):
-        if image:
-            return "", None, history + [[f"Image: {image}", None]]
-        else:
-            return "", None, history + [[user_message, None]]
+    full_response = ""
+    for partial_response in call_llama_vision_api(message, image_url):
+        full_response = partial_response
+        yield "", history + [(message, full_response)]
 
-    def bot(history):
-        user_message = history[-1][0]
-        image = None
-        if user_message.startswith("Image: "):
-            image = user_message.split(": ", 1)[1]
-            user_message = "What's in this image?"
-        
-        bot_message = chat_with_image(user_message, image, history[:-1])
-        history[-1][1] = ""
-        for character in bot_message:
-            history[-1][1] += character
-            yield history
+    history.append((message, full_response))
+    return "", history
 
-    msg.submit(user, [msg, image, chatbot], [msg, image, chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    )
+# Define the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Llama 3.2 Vision Chatbot Demo")
+    gr.Markdown("Enter your message and an image URL to analyze using the Llama 3.2 Vision model.")
+    
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Your message")
+    image_url = gr.Textbox(label="Image URL", value="https://napkinsdev.s3.us-east-1.amazonaws.com/next-s3-uploads/d96a3145-472d-423a-8b79-bca3ad7978dd/trello-board.png")
+    
+    clear = gr.Button("Clear")
+    
+    msg.submit(chat, [msg, chatbot, image_url], [msg, chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
-demo.queue()
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()