Update app.py

app.py

@@ -5,21 +5,14 @@ import base64
 from io import BytesIO
 from PIL import Image
 import numpy as np
-import json
 
 # Initialize the Together client
 api_key = os.environ.get('TOGETHER_API_KEY')
-client = None
-
-if api_key:
-    try:
-        client = Together(api_key=api_key)
-    except Exception as e:
-        print(f"Error initializing Together client: {e}")
+client = Together(api_key=api_key)
 
 def generate_gradio_app(image):
-    if not api_key or not client:
-        return "Error: TOGETHER_API_KEY not set or client initialization failed. Please check your API key."
+    if not api_key:
+        return "Error: TOGETHER_API_KEY not set. Please check your API key."
 
     try:
         # Convert numpy array to PIL Image
@@ -31,65 +24,53 @@ def generate_gradio_app(image):
         image.save(buffered, format="PNG")
         img_str = base64.b64encode(buffered.getvalue()).decode()
         
-        # Prepare the messages for the API call
-        system_message = "You are an AI assistant that can analyze wireframe images and generate detailed Gradio code based on their content."
-        user_message = f"""
-        <image>
-        data:image/png;base64,{img_str}
-        </image>
-        
-        Analyze this wireframe image and generate a complete Python code using Gradio that recreates all the main elements seen in the image. Provide a complete, runnable Gradio application.
-        """
-        
-        messages = [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": user_message}
-        ]
-        
+        # Prepare the message for the API call
+        # Using the format: <image>{base64_image}</image>\nHuman: {prompt}\nAssistant:
+        message = f"""<image>{img_str}</image>
+Human: Analyze this wireframe image and suggest a simple Gradio app layout based on it. Describe the main elements you see and how they could be implemented using Gradio components.
+Assistant: Certainly! I'll analyze the wireframe image and suggest a simple Gradio app layout based on what I see. Let me describe the main elements and how they could be implemented using Gradio components."""
+
         # Make the API call
         response = client.chat.completions.create(
             model="meta-llama/Llama-Vision-Free",
-            messages=messages,
-            max_tokens=5000,
+            messages=[{"role": "user", "content": message}],
+            max_tokens=512,
             temperature=0.7,
             top_p=0.7,
             top_k=50,
             repetition_penalty=1,
+            stop=["<|eot_id|>", "<|eom_id|>"],
+            stream=True
         )
         
-        # Extract the generated code from the response
-        generated_code = response.choices[0].message.content
-        return generated_code
+        # Collect the streamed response
+        generated_text = ""
+        for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                generated_text += chunk.choices[0].delta.content
+        
+        return generated_text
 
     except Exception as e:
         error_message = str(e)
-        if "400" in error_message:
-            return "Error 400: Bad Request. The API request was invalid. Please check your input and try again."
-        elif "401" in error_message:
-            return "Error 401: Unauthorized. Please check your API key and ensure it's correctly set."
-        elif "429" in error_message:
-            return "Error 429: Too Many Requests. Please wait a moment and try again."
-        elif "500" in error_message:
-            return "Error 500: Internal Server Error. There's an issue with the API service. Please try again later."
-        else:
-            return f"An unexpected error occurred: {error_message}\nPlease try again or contact support if the issue persists."
+        return f"An error occurred: {error_message}\nPlease try again or check your API key and connection."
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Turn your wireframe into a Gradio app")
-    gr.Markdown("Upload an image of your UI design and we'll build a Gradio app for you.")
+    gr.Markdown("# Analyze wireframe and suggest Gradio app layout")
+    gr.Markdown("Upload an image of your UI design for analysis and suggestions.")
     
     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 app", variant="primary")
+            generate_button = gr.Button("Analyze and Suggest", variant="primary")
         
         with gr.Column(scale=2):
-            code_output = gr.Code(language="python", label="Generated Gradio Code", lines=30)
+            text_output = gr.Textbox(label="Analysis and Suggestions", lines=10)
     
     generate_button.click(
         fn=generate_gradio_app,
         inputs=[image_input],
-        outputs=[code_output]
+        outputs=[text_output]
     )
 
 demo.launch()