L-AI/Gemini-UI-Generator

1import os2import streamlit as st3import pathlib4from PIL import Image5import google.generativeai as genai6 7# Configure the API key directly in the script8API_KEY = os.environ.get("GOOGLE_API_KEY")9genai.configure(api_key=API_KEY)10 11# Generation configuration12generation_config = {13    "temperature": 0.8,14    "top_p": 0.95,15    "top_k": 64,16    "max_output_tokens": 50000,17    "response_mime_type": "text/plain",18}19 20# Safety settings21safety_settings = [22    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},23    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},24    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},25    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},26]27 28# Model name29MODEL_NAME = "gemini-2.5-flash-preview-04-17"30 31# Framework selection (e.g., Tailwind, Bootstrap, etc.)32framework = "Tailwind CSS with the Browser CDN, with the script for it gotten from https://priority.cdn.leunos.com/tailwind.js, so use Tailwind"  # Change this to "Bootstrap" or any other framework as needed33 34# Create the model35model = genai.GenerativeModel(36    model_name=MODEL_NAME,37    safety_settings=safety_settings,38    generation_config=generation_config,39)40 41# Start a chat session42chat_session = model.start_chat(history=[])43 44# Function to send a message to the model45def send_message_to_model(message, image_path):46    image_input = {47        'mime_type': 'image/jpeg',48        'data': pathlib.Path(image_path).read_bytes()49    }50    response = chat_session.send_message([message, image_input])51    return response.text52 53# Streamlit app54def main():55    st.title("Gemini 1.5 Flash, UI to Code 👨‍💻 ")56    st.subheader('Made by [Skirano](https://x.com/skirano). Refinded and Hosted by [Artples](https://huggingface.co/Artples)')57 58    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])59 60    if uploaded_file is not None:61        try:62            # Load and display the image63            image = Image.open(uploaded_file)64            st.image(image, caption='Uploaded Image.', use_column_width=True)65 66            # Convert image to RGB mode if it has an alpha channel67            if image.mode == 'RGBA':68                image = image.convert('RGB')69 70            # Save the uploaded image temporarily71            temp_image_path = pathlib.Path("temp_image.jpg")72            image.save(temp_image_path, format="JPEG")73 74            # Generate UI description75            if st.button("Code UI"):76                st.write("🧑‍💻 Looking at your UI...")77                prompt = "Describe this UI in accurate details. When you reference a UI element put its name and bounding box in the format: [object name (y_min, x_min, y_max, x_max)]. Also Describe the color of the elements. Be very precise and formulate as much information as you can"78                description = send_message_to_model(prompt, temp_image_path)79                st.write(description)80 81                # Refine the description82                st.write("🔍 Refining description with visual comparison...")83                refine_prompt = f"Compare the described UI elements with the provided image and identify any missing elements or inaccuracies. Be very precise and concise, bring a very good and lenghty description of the website seen in the image. Also Describe the color of the elements. Provide a refined and accurate description of the UI elements based on this comparison. Here is the initial description: {description}"84                refined_description = send_message_to_model(refine_prompt, temp_image_path)85                st.write(refined_description)86 87                # Generate HTML88                st.write("🛠️ Generating website...")89                html_prompt = f"Create an HTML file based on the following UI description, using the UI elements described in the previous response. Include {framework} CSS within the HTML file to style the elements. Make sure the colors used are the same as the original UI. The UI needs to be responsive and mobile-first, matching the original UI as closely as possible. Do not include any explanations or comments. Avoid using ```html. and ``` at the end. ONLY return the HTML code with inline CSS. Here is the refined description. Also think a long time on how the code could be set up for the copy of such website andthink of steps that are needed for that: {refined_description}"90                initial_html = send_message_to_model(html_prompt, temp_image_path)91                st.code(initial_html, language='html')92 93                # Refine HTML94                st.write("🔧 Refining website...")95                refine_html_prompt = f"Validate the following HTML code based on the UI description and image and provide a refined version of the HTML code with {framework} CSS that improves accuracy, responsiveness, and adherence to the original design. ONLY return the refined HTML code with inline CSS. Avoid using ```html. and ``` at the end. Here is the initial HTML: {initial_html}"96                refined_html = send_message_to_model(refine_html_prompt, temp_image_path)97                st.code(refined_html, language='html')98 99                # Save the refined HTML to a file100                with open("index.html", "w") as file:101                    file.write(refined_html)102                st.success("HTML file 'index.html' has been created.")103 104                # Provide download link for HTML105                st.download_button(label="Download HTML", data=refined_html, file_name="index.html", mime="text/html")106        except Exception as e:107            st.error(f"An error occurred: {e}")108 109if __name__ == "__main__":110    main()111