import streamlit as st import pathlib from PIL import Image import google.generativeai as genai import os import shutil # --- Constants & Configuration --- # Use a more robust way to handle the API key API_KEY = st.secrets.get("GOOGLE_API_KEY") or os.getenv("GOOGLE_API_KEY") if not API_KEY: st.error("Google API Key not found. Please set it in Streamlit secrets or as an environment variable (`GOOGLE_API_KEY`).") st.stop() genai.configure(api_key=API_KEY) GENERATION_CONFIG = { "temperature": 0.8, # Slightly lower temperature for more predictable code "top_p": 0.9, "top_k": 40, "max_output_tokens": 8192, "response_mime_type": "text/plain", } 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 = "gemini-1.5-pro-latest" # --- Model & Chat Session Management --- @st.cache_resource def get_gemini_model(): """Loads and caches the Google GenerativeModel.""" return genai.GenerativeModel( model_name=MODEL_NAME, safety_settings=SAFETY_SETTINGS, generation_config=GENERATION_CONFIG, ) def get_chat_session(): """Initializes and returns a chat session, caching it in Streamlit's session state.""" if "chat_session" not in st.session_state: st.session_state.chat_session = get_gemini_model().start_chat(history=[]) return st.session_state.chat_session # --- Helper Function for Model Communication --- def send_message_with_image(chat_session, message: str, image_path: pathlib.Path) -> str: """Sends a message and an image to the Gemini model.""" try: response = chat_session.send_message([message, genai.upload_file(image_path)]) return response.text except Exception as e: st.error(f"Error communicating with the Gemini model: {e}") st.exception(e) 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", initial_sidebar_state="collapsed") st.title("Gemini 1.5 Pro: Images to Code 👨‍💻") st.markdown('Made with ❤️ by [KhulnaSoft](https://x.com/khulnasoft) and Enhanced by Google') st.info("Upload an image of a UI design, and I'll generate the corresponding HTML and CSS code for you!") # Framework selection framework_options = { "Regular CSS (Flexbox/Grid)": "Regular CSS with flexbox and grid layouts", "Bootstrap": "Bootstrap framework with its utility classes", "Tailwind CSS": "Tailwind CSS with its utility classes", "Materialize CSS": "Materialize CSS framework" } selected_framework = st.selectbox( "Choose your preferred CSS framework:", options=list(framework_options.keys()), help="This will influence the CSS generated within your HTML file." ) uploaded_file = st.file_uploader("Upload a UI image (JPG, JPEG, PNG):", type=["jpg", "jpeg", "png"]) # Use a temporary directory for file handling temp_dir = pathlib.Path("temp_uploads") if not temp_dir.exists(): temp_dir.mkdir() temp_image_path = temp_dir / "uploaded_image.jpg" if uploaded_file: try: # Load, display, and save the image image = Image.open(uploaded_file).convert('RGB') st.image(image, caption='Uploaded UI Image', use_column_width=True) image.save(temp_image_path, format="JPEG") st.markdown("---") # Use st.form for a better user experience with a single button with st.form("code_generation_form"): st.subheader("Generate UI Code") st.markdown(f"**Selected Framework:** `{selected_framework}`") submitted = st.form_submit_button("Generate Code") if submitted: # Use a single progress bar for all steps progress_bar = st.progress(0, text="Starting code generation...") chat_session = get_chat_session() framework_instruction = framework_options[selected_framework] # Step 1: Generate & Refine Description progress_bar.progress(25, text="Step 1/2: Analyzing UI and generating description...") prompt_description = ( f"Describe this UI in accurate details, including UI elements, their bounding boxes in the format: " f"[object name (y_min, x_min, y_max, x_max)], and their colors. " f"Then, refine this description by comparing it to the image for any missing elements or inaccuracies. " f"Provide a final, refined, and accurate description based on this analysis." ) refined_description = send_message_with_image(chat_session, prompt_description, temp_image_path) if not refined_description.startswith("An error occurred"): st.success("UI Description Generated!") with st.expander("See Generated UI Description"): st.text(refined_description) else: st.error("Failed to generate UI description.") return # Step 2: Generate & Refine HTML progress_bar.progress(75, text="Step 2/2: Generating and refining HTML/CSS code...") prompt_html = ( f"Create a single HTML file based on the following UI description: '{refined_description}'. " f"Use {framework_instruction} to style the elements. The UI must be responsive and " f"mobile-first, matching the original design as closely as possible. " f"Validate the generated code and provide a refined version that improves accuracy, " f"responsiveness, and adherence to the original design. " f"ONLY return the complete, refined HTML code. Do not include any explanations, comments, or ```html block." ) refined_html = send_message_with_image(chat_session, prompt_html, temp_image_path) if not refined_html.startswith("An error occurred"): st.success("HTML Refined Successfully!") st.subheader("Refined Generated HTML:") st.code(refined_html, language='html') st.markdown("---") st.success("All steps completed! Your `index.html` file is ready for download.") # Provide download link 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.") else: st.error("Failed to generate HTML code.") progress_bar.empty() # Hide the progress bar except Exception as e: st.error(f"An unexpected error occurred: {e}") st.exception(e) finally: if temp_dir.exists(): shutil.rmtree(temp_dir) # Use shutil to remove the directory and its contents else: st.write("Please upload an image to start generating UI code.") if __name__ == "__main__": main()