Update app.py

app.py

@@ -11,6 +11,7 @@ import base64
 import uuid
 import concurrent.futures
 import re
+import tempfile
 
 # --- Configuration and Initialization ---
 
@@ -87,7 +88,7 @@ Based *only* on the requirements above, generate a single, comprehensive JSON ob
 """
 
 def generate_image(prompt: str, aspect_ratio: str):
-    """Generates an image using Gemini and returns a base64 data URL."""
+    """Generates an image using Gemini and returns a file path and a base64 data URL."""
     print(f"Generating image for prompt: {prompt}")
     try:
         # Refactored to use the new client.models.generate_content method
@@ -103,14 +104,25 @@ def generate_image(prompt: str, aspect_ratio: str):
 
         if image_part:
             mime_type = image_part.inline_data.mime_type
-            data = base64.b64encode(image_part.inline_data.data).decode("utf-8")
-            return f"data:{mime_type};base64,{data}"
+            image_data = image_part.inline_data.data
+
+            # Create base64 URL
+            base64_data = base64.b64encode(image_data).decode("utf-8")
+            base64_url = f"data:{mime_type};base64,{base64_data}"
+
+            # Save to temp file
+            extension = mime_type.split('/')[-1]
+            with tempfile.NamedTemporaryFile(delete=False, suffix=f".{extension}") as temp_file:
+                temp_file.write(image_data)
+                temp_filepath = temp_file.name
+            
+            return temp_filepath, base64_url
         else:
             print(f"Warning: Image generation failed for prompt: {prompt}. No image data in response.")
-            return ""
+            return None, None
     except Exception as e:
         print(f"Error generating image for prompt '{prompt}': {e}")
-        return ""
+        return None, None
 
 def generate_page_html(page, settings):
     """Generates the HTML for a single page based on the website settings."""
@@ -388,7 +400,7 @@ def create_gradio_app():
         
         # --- Step 2: Generate Images in Parallel ---
         tasks = []
-        logo_prompts = image_prompts.get('logoPrompts', [])
+        logo_prompts = [re.sub(r'import\s.*', '', p).strip() for p in image_prompts.get('logoPrompts', [])]
         block_prompts = image_prompts.get('blockImagePrompts', [])
         
         if logo_prompts:
@@ -399,20 +411,23 @@ def create_gradio_app():
         if block_prompts:
             yield {status_text: f"Generating {len(block_prompts)} content images..."}
             for p_info in block_prompts:
+                p_info['prompt'] = re.sub(r'import\s.*', '', p_info.get('prompt', '')).strip()
                 tasks.append(("block", p_info))
         
         with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
             future_to_task = {executor.submit(generate_image, task[1]['prompt'] if task[0]=='block' else task[1], '16:9' if task[0]=='block' and task[1].get('type')=='hero' else '1:1'): task for task in tasks}
             
-            logo_results = []
+            logo_filepaths = []
+            logo_base64s = []
             for i, future in enumerate(concurrent.futures.as_completed(future_to_task)):
                 task_type, task_data = future_to_task[future]
                 try:
-                    image_url = future.result()
-                    if not image_url: continue
+                    filepath, base64_url = future.result()
+                    if not filepath or not base64_url: continue
 
                     if task_type == "logo":
-                        logo_results.append(image_url)
+                        logo_filepaths.append(filepath)
+                        logo_base64s.append(base64_url)
                     
                     elif task_type == "block":
                         original_block_id = task_data['blockId']
@@ -427,7 +442,7 @@ def create_gradio_app():
                         if not block_to_update: continue
 
                         key_to_update = 'bgImage' if block_to_update.get('type') == 'hero' else 'image'
-                        block_to_update[key_to_update] = image_url
+                        block_to_update[key_to_update] = base64_url
                         
                     yield {status_text: f"Processing images... ({i+1}/{len(tasks)})"}
 
@@ -437,13 +452,15 @@ def create_gradio_app():
         yield {
             generating_view: gr.update(visible=False),
             logo_picker_view: gr.update(visible=True),
-            logo_gallery: gr.update(value=logo_results if logo_results else []),
+            logo_gallery: gr.update(value=logo_filepaths if logo_filepaths else []),
             website_settings_state: new_settings,
-            active_page_path_state: active_page_path
+            active_page_path_state: active_page_path,
+            generated_logos_state: logo_base64s
         }
         
-    def handle_logo_select(evt: gr.SelectData, settings, active_path):
-        logo_url = evt.value
+    def handle_logo_select(evt: gr.SelectData, settings, active_path, all_logo_base64s):
+        selected_index = evt.index
+        logo_url = all_logo_base64s[selected_index]
         settings['header']['logoUrl'] = logo_url
         
         page = get_page_by_path(active_path, settings)
@@ -532,6 +549,7 @@ def create_gradio_app():
         # --- State Management ---
         website_settings_state = gr.State(DEFAULT_WEBSITE_SETTINGS)
         active_page_path_state = gr.State("index.html")
+        generated_logos_state = gr.State([])
 
         # --- View: Welcome ---
         with gr.Column(visible=True) as welcome_view:
@@ -625,13 +643,14 @@ def create_gradio_app():
             inputs=questionnaire_inputs,
             outputs=[
                 questionnaire_view, generating_view, status_text, error_box,
-                logo_picker_view, logo_gallery, website_settings_state, active_page_path_state
+                logo_picker_view, logo_gallery, website_settings_state, active_page_path_state,
+                generated_logos_state
             ]
         )
         
         logo_gallery.select(
             handle_logo_select,
-            inputs=[website_settings_state, active_page_path_state],
+            inputs=[website_settings_state, active_page_path_state, generated_logos_state],
             outputs=[
                 logo_picker_view, preview_view, website_settings_state, html_preview, code_preview,
                 page_selector, header_title_input, footer_text_input,