import streamlit as st import pathlib from PIL import Image import google.generativeai as genai import os # Import os for environment variables # --- Configuration --- try: # Ensure GOOGLE_API_KEY is set in Streamlit secrets or as an environment variable API_KEY = st.secrets["GOOGLE_API_KEY"] except KeyError: st.error("Google API Key not found. Please set it in Streamlit secrets or as an environment variable (`GOOGLE_API_KEY`).") st.stop() # Stop the app if API key is missing genai.configure(api_key=API_KEY) # Generation configuration for the Gemini model GENERATION_CONFIG = { "temperature": 1, "top_p": 0.95, "top_k": 64, "max_output_tokens": 8192, "response_mime_type": "text/plain", } # Safety settings for the Gemini model SAFETY_SETTINGS = [ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}, {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}, {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"}, {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"}, ] # Model name constant MODEL_NAME = "gemini-1.5-pro-latest" # --- Model Initialization --- @st.cache_resource # Cache the model to avoid re-initializing on every rerun def load_gemini_model(): """Loads and caches the Google GenerativeModel.""" return genai.GenerativeModel( model_name=MODEL_NAME, safety_settings=SAFETY_SETTINGS, generation_config=GENERATION_CONFIG, ) model = load_gemini_model() # Initialize chat session in Streamlit's session state # This ensures the chat history persists across reruns for a single user if "chat_session" not in st.session_state: st.session_state.chat_session = model.start_chat(history=[]) # --- Helper Function for Model Communication --- def send_message_to_model(message: str, image_path: pathlib.Path) -> str: """Sends a message and an image to the Gemini model and returns the response.""" image_input = { 'mime_type': 'image/jpeg', 'data': image_path.read_bytes() } try: response = st.session_state.chat_session.send_message([message, image_input]) return response.text except Exception as e: st.error(f"Error communicating with the Gemini model: {e}") st.exception(e) # Display full traceback for debugging return "An error occurred during AI model communication. Please try again or check your API key." # --- Streamlit App --- def main(): """Main function to run the Streamlit application.""" st.set_page_config(page_title="Gemini 1.5 Pro: Images to Code", layout="wide") # Set a wider layout st.title("Gemini 1.5 Pro: Images to Code 👨‍💻") st.markdown('Made with ❤️ by [KhulnaSoft](https://x.com/khulnasoft)') st.info("Upload an image of a UI design, and I'll generate the corresponding HTML and CSS code for you!") # Framework selection using a selectbox framework_options = { "Regular CSS (Flexbox/Grid)": "Regular CSS use flex grid etc", "Bootstrap": "Bootstrap", "Tailwind CSS": "Tailwind CSS", "Materialize CSS": "Materialize CSS" } selected_framework_name = st.selectbox( "Choose your preferred CSS framework:", options=list(framework_options.keys()), help="This will influence the CSS generated within your HTML file." ) framework = framework_options[selected_framework_name] uploaded_file = st.file_uploader("Upload a UI image (JPG, JPEG, PNG):", type=["jpg", "jpeg", "png"]) # temp_image_path is declared outside the try block to ensure it's accessible for cleanup temp_image_path = pathlib.Path("temp_image.jpg") if uploaded_file is not None: try: # Load and display the image image = Image.open(uploaded_file) st.image(image, caption='Uploaded UI Image', use_column_width=True) # Convert image to RGB mode if it has an alpha channel if image.mode == 'RGBA': image = image.convert('RGB') # Save the uploaded image temporarily image.save(temp_image_path, format="JPEG") st.markdown("---") # Visual separator # Button to trigger the generation process if st.button("Generate UI Code", help="Click to initiate the multi-step code generation."): st.subheader("Code Generation Process:") # Step 1: Generate initial UI description with st.spinner("Step 1/4: Describing your UI elements and colors..."): 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." description = send_message_to_model(prompt, temp_image_path) st.success("UI Description Generated!") with st.expander("See Initial UI Description"): st.text(description) # Step 2: Refine the description with st.spinner("Step 2/4: Refining description with visual comparison..."): refine_prompt = f"Compare the described UI elements with the provided image and identify any missing elements or inaccuracies. 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}" refined_description = send_message_to_model(refine_prompt, temp_image_path) st.success("UI Description Refined!") with st.expander("See Refined UI Description"): st.text(refined_description) # Step 3: Generate initial HTML with st.spinner("Step 3/4: Generating initial HTML with CSS..."): html_prompt = ( f"Create an HTML file based on the following UI description, using the UI elements described in the previous response. " f"Include {framework} CSS within the HTML file to style the elements. " f"Make sure the colors used are the same as the original UI. " f"The UI needs to be responsive and mobile-first, matching the original UI as closely as possible. " f"Do not include any explanations or comments. Avoid using ```html and ``` at the end. " f"ONLY return the HTML code with inline CSS. Here is the refined description: {refined_description}" ) initial_html = send_message_to_model(html_prompt, temp_image_path) st.success("Initial HTML Generated!") st.subheader("Initial Generated HTML:") st.code(initial_html, language='html') # Step 4: Refine HTML with st.spinner("Step 4/4: Refining the generated HTML code..."): 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. " f"ONLY return the refined HTML code with inline CSS. Avoid using ```html and ``` at the end. " f"Here is the initial HTML: {initial_html}" ) refined_html = send_message_to_model(refine_html_prompt, temp_image_path) st.success("HTML Refined Successfully!") st.subheader("Refined Generated HTML:") st.code(refined_html, language='html') st.markdown("---") # Final separator st.success("All steps completed! Your `index.html` file is ready for download.") # Save the refined HTML to a file and provide download link with open("index.html", "w", encoding="utf-8") as file: # Specify encoding file.write(refined_html) st.download_button( label="Download index.html", data=refined_html, file_name="index.html", mime="text/html" ) st.info("You can open the downloaded `index.html` file in your web browser to view the generated UI.") except Exception as e: st.error(f"An unexpected error occurred: {e}") st.exception(e) # Displays the full traceback finally: # Clean up the temporary image file whether an error occurred or not if temp_image_path.exists(): os.remove(temp_image_path) # st.success("Temporary image file removed.") # Can uncomment for debugging else: st.write("Please upload an image to start generating UI code.") if __name__ == "__main__": main()