Update app.py

app.py

@@ -105,21 +105,24 @@ def gen(prompts, width, height, model_name, num_variants=8):
 def display_gallery():
     st.header("Galería de Imágenes Guardadas")
     images = list_saved_images()
+    image_options = {str(image_file): f"Imagen {i+1}" for i, image_file in enumerate(images)}
+    selected_image = st.radio("Seleccionar imagen para FaceSwap", options=image_options.keys(), index=None)    
+    if selected_image:
+        st.session_state['generated_image_path'] = selected_image
+        prompt = get_prompt_for_image(Path(selected_image).name)
+        st.write(f"Prompt: {prompt[:100]}")
+    if selected_image and st.button("Aplicar FaceSwap"):
+        st.success("Imagen seleccionada para FaceSwap")
     if images:
-        cols = st.columns(8)
+        cols = st.columns(4) 
         for i, image_file in enumerate(images):
-            with cols[i % 8]:
-                st.image(str(image_file), caption=image_file.name, use_column_width=True)
-                prompt = get_prompt_for_image(image_file.name)
-                st.write(prompt[:100])
-                if st.button(f"FaceSwap", key=f"select_{i}_{image_file.name}"):
-                    st.session_state['generated_image_path'] = str(image_file)
-                    st.success("Imagen seleccionada")
-                if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"):
+            with cols[i % 4]:
+                st.image(str(image_file), use_column_width=True)
+                if st.button(f"Borrar", key=f"delete_{i}_{image_file}"):
                     if image_file.exists():
                         os.remove(image_file)
                         st.success("Imagen borrada")
-                        display_gallery()
+                        st.rerun()
 def get_prompt_for_image(image_name):
     try:
         with open(DATA_PATH / "prompts.txt", "r") as f: