Update app.py

app.py

@@ -1,30 +1,47 @@
 import gradio as gr
+from PIL import Image
+import requests
+import os
 from together import Together
 import base64
 from io import BytesIO
-from PIL import Image
 import numpy as np
 import traceback
-import json
+
+# Initialize Together client
+client = None
+
+def initialize_client(api_key=None):
+    global client
+    if api_key:
+        client = Together(api_key=api_key)
+    elif "TOGETHER_API_KEY" in os.environ:
+        client = Together()
+    else:
+        raise ValueError("Please provide an API key or set the TOGETHER_API_KEY environment variable")
+
+def encode_image(image, max_size=(800, 800), quality=85):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
+    
+    image.thumbnail(max_size)
+    if image.mode in ('RGBA', 'LA'):
+        background = Image.new(image.mode[:-1], image.size, (255, 255, 255))
+        background.paste(image, mask=image.split()[-1])
+        image = background
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG", quality=quality)
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
 def generate_gradio_app(api_key, image):
     if not api_key:
         return "Error: API key not provided. Please enter your Together API key."
 
     try:
-        # Initialize the Together client with the provided API key
-        client = Together(api_key=api_key)
+        initialize_client(api_key)
 
-        # Convert numpy array to PIL Image
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image.astype('uint8'), 'RGB')
+        encoded_image = encode_image(image)
         
-        # 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:
 
 1. Briefly describe the main elements of the UI.
@@ -39,58 +56,39 @@ def generate_gradio_app(api_key, image):
 
 Please generate the Gradio code based on the provided image, focusing on the most crucial elements to fit within the token limit."""
 
-        # Make the API call with the corrected message format
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}}
+                ]
+            }
+        ]
+
         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>"
-                }
-            ],
+            messages=messages,
             max_tokens=2000,
             temperature=0.7,
             top_p=0.7,
             top_k=50,
             repetition_penalty=1,
-            stop=["<|eot_id|>", "<|eom_id|>"],
-            stream=False
+            stop=["<|eot_id|>", "<|eom_id|>"]
         )
-        
-        # Debug: Print relevant parts of the response
-        print("API Response received. Content available:", bool(response.choices))
 
-        # 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))
+            return 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
 
     except Exception as e:
         error_message = str(e)
         stack_trace = traceback.format_exc()
-        
-        # Check if the error is related to the API response
-        if "TogetherErrorResponse" in error_message:
-            try:
-                # Try to parse the error message as JSON
-                error_data = json.loads(error_message.split("TogetherErrorResponse")[-1].strip())
-                if isinstance(error_data.get('code'), int):
-                    error_data['code'] = str(error_data['code'])  # Convert code to string
-                error_message = f"API Error: {error_data.get('message', 'Unknown error')}"
-            except json.JSONDecodeError:
-                # If parsing fails, use the original error message
-                pass
-        
         return f"An error occurred: {error_message}\n\nStack trace:\n{stack_trace}\n\nPlease check your API key and try again."
 
-# The rest of your code remains the same
 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.")