|
import gradio as gr |
|
import os |
|
from together import Together |
|
import base64 |
|
from io import BytesIO |
|
from PIL import Image |
|
import numpy as np |
|
|
|
|
|
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}") |
|
|
|
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." |
|
|
|
try: |
|
|
|
if isinstance(image, np.ndarray): |
|
image = Image.fromarray(image.astype('uint8'), 'RGB') |
|
|
|
|
|
buffered = BytesIO() |
|
image.save(buffered, format="PNG") |
|
img_str = base64.b64encode(buffered.getvalue()).decode() |
|
|
|
|
|
system_message = "You are an AI assistant that can analyze wireframe images and generate Gradio code based on their content." |
|
user_message = f""" |
|
<image> |
|
data:image/png;base64,{img_str} |
|
</image> |
|
|
|
Analyze this wireframe image and generate Python code using Gradio that could recreate the main elements seen in the image. Use Gradio components that best represent the UI elements in the wireframe. |
|
""" |
|
|
|
messages = [ |
|
{"role": "system", "content": system_message}, |
|
{"role": "user", "content": user_message} |
|
] |
|
|
|
|
|
response = client.chat.completions.create( |
|
model="meta-llama/Llama-Vision-Free", |
|
messages=messages, |
|
max_tokens=50000, |
|
temperature=0.7, |
|
top_p=0.7, |
|
top_k=50, |
|
repetition_penalty=1, |
|
stop=["<|eot_id|>", "<|eom_id|>"], |
|
stream=True |
|
) |
|
|
|
|
|
generated_code = "" |
|
for chunk in response: |
|
if chunk.choices[0].delta.content is not None: |
|
generated_code += chunk.choices[0].delta.content |
|
|
|
return generated_code |
|
except Exception as e: |
|
return f"An error occurred: {str(e)}" |
|
|
|
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.") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
image_input = gr.Image(label="Upload a screenshot", elem_id="image_upload") |
|
example_link = gr.Markdown("Need an example image? [Try ours](https://example.com/wireframe.png).") |
|
|
|
model_dropdown = gr.Dropdown( |
|
choices=["Llama-Vision-Free"], |
|
value="Llama-Vision-Free", |
|
label="AI Model" |
|
) |
|
|
|
generate_button = gr.Button("Generate Gradio app", variant="primary") |
|
|
|
code_output = gr.Code(language="python", label="Generated Gradio Code", lines=20) |
|
|
|
generate_button.click( |
|
fn=generate_gradio_app, |
|
inputs=[image_input], |
|
outputs=[code_output] |
|
) |
|
|
|
demo.launch() |