Update app.py

app.py

@@ -1,100 +1,85 @@
+import os
 import gradio as gr
 from together import Together
-import base64
-from io import BytesIO
 from PIL import Image
-import numpy as np
-import traceback
-
-def generate_gradio_app(api_key, image):
-    if not api_key:
-        return "Error: API key not provided. Please enter your Together API key."
+import io
+import base64
 
-    try:
-        # Initialize the Together client with the provided API key
-        client = Together(api_key=api_key)
+# Initialize the Together AI client
+client = Together(api_key=os.environ.get('TOGETHER_API_KEY'))
 
-        # Convert numpy array to PIL Image
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image.astype('uint8'), 'RGB')
-        
-        # Convert the image to base64
-        buffered = BytesIO()
-        image.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode()
-        
-        # Prepare the prompt
-        prompt = """You are an AI assistant specialized in UI/UX design and Gradio app development. Analyze the attached screenshot or UI mockup and generate a concise Gradio code that recreates the main elements of this design. Follow these steps:
+def encode_image(image):
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-1. Briefly describe the main elements of the UI.
-2. Generate a compact Gradio Python code that recreates the essential UI elements.
-3. Use appropriate Gradio components for key elements in the UI.
-4. Include necessary imports at the beginning of the code.
-5. Implement minimal placeholder functions for interactive elements.
-6. Use gr.Blocks() to create a basic layout that captures the essence of the screenshot.
-7. Include the gr.Blocks().launch() call at the end of the code.
-8. Provide a runnable Gradio application focusing on the most important aspects of the UI.
-9. Keep the code concise, aiming for no more than 2000 tokens.
+def chat_with_image(message, image, history):
+    # Encode the image
+    if image is not None:
+        encoded_image = encode_image(Image.open(image))
+        image_message = {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded_image}"}},
+                {"type": "text", "text": message}
+            ]
+        }
+    else:
+        image_message = {"role": "user", "content": message}
+    
+    # Prepare the messages
+    messages = [{"role": "system", "content": "You are a helpful assistant."}]
+    for human, assistant in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": assistant})
+    messages.append(image_message)
+    
+    # Call the Together AI API
+    response = client.chat.completions.create(
+        model="meta-llama/Llama-Vision-Free",
+        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
 
-Please generate the Gradio code based on the provided image, focusing on the most crucial elements to fit within the token limit."""
+    return full_response
 
-        # Make the API call with the corrected message format
-        response = client.chat.completions.create(
-            model="meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
-            messages=[
-                {
-                    "role": "user",
-                    "content": prompt + f"\n\n<image>\ndata:image/png;base64,{img_str}\n</image>"
-                }
-            ],
-            max_tokens=2000,
-            temperature=0.7,
-            top_p=0.7,
-            top_k=50,
-            repetition_penalty=1,
-            stop=["<|eot_id|>", "<|eom_id|>"],
-            stream=False
-        )
-        
-        # Debug: Print relevant parts of the response
-        print("API Response received. Content available:", bool(response.choices))
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    image = gr.Image(type="filepath")
+    clear = gr.Button("Clear")
 
-        # Extract the generated code from the response
-        if response.choices and response.choices[0].message:
-            generated_code = response.choices[0].message.content
-            print("Generated code length:", len(generated_code))
-        else:
-            return "Error: Unexpected response structure from the API."
-        
-        if not generated_code:
-            return "Error: No code generated from the model. Please try again."
-        
-        return generated_code
+    def user(user_message, image, history):
+        return "", image, history + [[user_message, None]]
 
-    except Exception as e:
-        error_message = str(e)
-        stack_trace = traceback.format_exc()
-        return f"An error occurred: {error_message}\n\nStack trace:\n{stack_trace}\n\nPlease check your API key and try again."
+    def bot(history):
+        user_message, image = history[-1][0], None
+        if len(history) > 1 and isinstance(history[-2][0], dict):
+            image = history[-2][0]['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
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Generate Concise Gradio App from Wireframe")
-    gr.Markdown("Enter your Together API key, upload an image of your UI design, and we'll generate a compact Gradio code to recreate its main elements.")
-    
-    api_key_input = gr.Textbox(label="Enter your Together API Key", type="password")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(label="Upload a screenshot", elem_id="image_upload")
-            generate_button = gr.Button("Generate Gradio Code", variant="primary")
-        
-        with gr.Column(scale=2):
-            code_output = gr.Code(language="python", label="Generated Gradio Code", lines=30)
-    
-    generate_button.click(
-        fn=generate_gradio_app,
-        inputs=[api_key_input, image_input],
-        outputs=[code_output]
+    msg.submit(user, [msg, image, chatbot], [msg, image, chatbot], queue=False).then(
+        bot, chatbot, chatbot
     )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
-if __name__ == "__main__":
-    demo.launch(debug=True)
+demo.queue()
+demo.launch()