GeminiCode / app.py
khulnasoft's picture
Update app.py
ced0116 verified
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()