Update app.py

app.py

@@ -2,28 +2,27 @@ import streamlit as st
 import pathlib
 from PIL import Image
 import google.generativeai as genai
-import os # Import os for environment variables
+import os
+import shutil
 
-# --- 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:
+# --- 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() # Stop the app if API key is missing
+    st.stop()
 
 genai.configure(api_key=API_KEY)
 
-# Generation configuration for the Gemini model
 GENERATION_CONFIG = {
-    "temperature": 1,
-    "top_p": 0.95,
-    "top_k": 64,
+    "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 for the Gemini model
 SAFETY_SETTINGS = [
     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
@@ -31,12 +30,11 @@ SAFETY_SETTINGS = [
     {"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():
+# --- Model & Chat Session Management ---
+@st.cache_resource
+def get_gemini_model():
     """Loads and caches the Google GenerativeModel."""
     return genai.GenerativeModel(
         model_name=MODEL_NAME,
@@ -44,142 +42,136 @@ def load_gemini_model():
         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=[])
+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_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()
-    }
+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 = st.session_state.chat_session.send_message([message, image_input])
+        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) # Display full traceback for debugging
+        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") # Set a wider layout
+    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)')
+    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 using a selectbox
+    # Framework selection
     framework_options = {
-        "Regular CSS (Flexbox/Grid)": "Regular CSS use flex grid etc",
-        "Bootstrap": "Bootstrap",
-        "Tailwind CSS": "Tailwind CSS",
-        "Materialize CSS": "Materialize CSS"
+        "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_name = st.selectbox(
+    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."
     )
-    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:
+    
+    # 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 and display the image
-            image = Image.open(uploaded_file)
+            # Load, display, and save the image
+            image = Image.open(uploaded_file).convert('RGB')
             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.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 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"):
+                    with st.expander("See Generated 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)
+                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("---") # 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.")
-
+                    
+                    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) # Displays the full traceback
-
+            st.exception(e)
+            
         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
+            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.")