Spaces:
Running
Running
| 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 --- | |
| 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() |