Complete Refactored Code

Improvement:

Constants: Define constants at the top for model names and safety settings to improve readability and maintainability.
Function for Image Saving: While a small task, encapsulating temporary file saving could be a good practice for larger applications. For this case, it's fine as is.
Redundant pathlib for temp_image_path: You create a Path object but then pass it to read_bytes() which expects a string or Path-like object. It's okay, but str(temp_image_path) could be used for clarity if needed by other functions, although read_bytes() handles Path directly.

Improvement: Streamlit reruns the script from top to bottom whenever there's an interaction. Re-initializing the GenerativeModel every time is inefficient.

Change: Use

@st
.cache_resource to cache the model.

Improvement: The chat_session should ideally be managed within st.session_state to maintain a persistent conversation with the model across Streamlit reruns for a single user, if you intend to have a continued conversation. While your current use case is step-by-step generation, it's a good practice for chat applications.

app.py

@@ -2,13 +2,20 @@ 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
 
-# Configure the API key directly in the script
-API_KEY = 'YOUR KEY'
 genai.configure(api_key=API_KEY)
 
-# Generation configuration
-generation_config = {
+# Generation configuration for the Gemini model
+GENERATION_CONFIG = {
     "temperature": 1,
     "top_p": 0.95,
     "top_k": 64,
@@ -16,94 +23,165 @@ generation_config = {
     "response_mime_type": "text/plain",
 }
 
-# Safety settings
-safety_settings = [
+# 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
+# Model name constant
 MODEL_NAME = "gemini-1.5-pro-latest"
 
-# Framework selection (e.g., Tailwind, Bootstrap, etc.)
-framework = "Regular CSS use flex grid etc"  # Change this to "Bootstrap" or any other framework as needed
-
-# Create the model
-model = genai.GenerativeModel(
-    model_name=MODEL_NAME,
-    safety_settings=safety_settings,
-    generation_config=generation_config,
-)
-
-# Start a chat session
-chat_session = model.start_chat(history=[])
-
-# Function to send a message to the model
-def send_message_to_model(message, image_path):
+# --- 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': pathlib.Path(image_path).read_bytes()
+        'data': image_path.read_bytes()
     }
-    response = chat_session.send_message([message, image_input])
-    return response.text
-
-# Streamlit app
+    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():
-    st.title("Gemini 1.5 Pro, Images to Code 👨‍💻 ")
-    st.subheader('Made with ❤️ by [KhulnaSoft](https://x.com/khulnasoft)')
+    """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("Choose an image...", type=["jpg", "jpeg", "png"])
+    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 Image.', use_column_width=True)
+            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
-            temp_image_path = pathlib.Path("temp_image.jpg")
             image.save(temp_image_path, format="JPEG")
 
-            # Generate UI description
-            if st.button("Code UI"):
-                st.write("🧑‍💻 Looking at your UI...")
-                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.write(description)
-
-                # Refine the description
-                st.write("🔍 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.write(refined_description)
-
-                # Generate HTML
-                st.write("🛠️ Generating website...")
-                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: {refined_description}"
-                initial_html = send_message_to_model(html_prompt, temp_image_path)
-                st.code(initial_html, language='html')
-
-                # Refine HTML
-                st.write("🔧 Refining website...")
-                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}"
-                refined_html = send_message_to_model(refine_html_prompt, temp_image_path)
-                st.code(refined_html, language='html')
-
-                # Save the refined HTML to a file
-                with open("index.html", "w") as file:
+            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.success("HTML file 'index.html' has been created.")
 
-                # Provide download link for HTML
-                st.download_button(label="Download HTML", data=refined_html, file_name="index.html", mime="text/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 error occurred: {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()