Update app.py

app.py

@@ -5,7 +5,6 @@ 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')
@@ -25,31 +24,41 @@ def generate_gradio_app(image):
         image.save(buffered, format="PNG")
         img_str = base64.b64encode(buffered.getvalue()).decode()
         
-        # Prepare the message for the API call
-        message_content = f"""<image>{img_str}</image>
-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."""
-
-        # Prepare the API request payload
-        payload = {
-            "model": "meta-llama/Llama-Vision-Free",
-            "messages": [{"role": "user", "content": message_content}],
-            "max_tokens": 512,
-            "temperature": 0.7,
-            "top_p": 0.7,
-            "top_k": 50,
-            "repetition_penalty": 1,
-            "stop": ["<|eot_id|>", "<|eom_id|>"],
-            "stream_tokens": True
-        }
+        # Prepare the prompt
+        prompt = """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 as a Gradio app, 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. After describing the UI, suggest how this could be implemented using Gradio components."""
 
         # Make the API call
-        response = client.chat.completions.create(**payload)
+        stream = client.chat.completions.create(
+            model="meta-llama/Llama-Vision-Free",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{img_str}",
+                            },
+                        },
+                    ],
+                }
+            ],
+            max_tokens=2048,
+            temperature=0.7,
+            top_p=0.7,
+            top_k=50,
+            repetition_penalty=1,
+            stop=["<|eot_id|>", "<|eom_id|>"],
+            stream=True
+        )
         
         # Collect the streamed response
         generated_text = ""
-        for chunk in response:
-            if chunk.choices and chunk.choices[0].delta.content:
+        for chunk in stream:
+            if chunk.choices[0].delta.content is not None:
                 generated_text += chunk.choices[0].delta.content
+                yield generated_text  # Yield partial results for gradio to update in real-time
         
         return generated_text
 
@@ -67,7 +76,7 @@ with gr.Blocks() as demo:
             generate_button = gr.Button("Analyze and Suggest", variant="primary")
         
         with gr.Column(scale=2):
-            text_output = gr.Textbox(label="Analysis and Suggestions", lines=10)
+            text_output = gr.Textbox(label="Analysis and Suggestions", lines=20)
     
     generate_button.click(
         fn=generate_gradio_app,