Update app.py

app.py

@@ -1,16 +1,21 @@
-import os, random, asyncio
+import os
+import random
+import asyncio
 from pathlib import Path
 from PIL import Image
 from insightface.app import FaceAnalysis
 import streamlit as st
 from huggingface_hub import InferenceClient, AsyncInferenceClient
 import yaml
-import insightface
 
-try: 
-    with open("config.yaml", "r") as file: credentials = yaml.safe_load(file)
-except Exception: credentials = {"username": "", "password": ""} 
+# Cargar credenciales
+try:
+    with open("config.yaml", "r") as file:
+        credentials = yaml.safe_load(file)
+except Exception:
+    credentials = {"username": "", "password": ""}
 
+# Configuraciones iniciales
 MAX_SEED = 2**31 - 1
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
 client = AsyncInferenceClient()
@@ -30,31 +35,37 @@ def run_async(func):
     return loop.run_until_complete(loop.run_in_executor(None, func))
 
 async def generate_image(prompt, model, width, height, scales, steps, seed):
-    if seed == -1: seed = random.randint(0, MAX_SEED)
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
     try:
         return await client.text_to_image(prompt=prompt, height=height, width=width, guidance_scale=scales, num_inference_steps=steps, model=model), seed
-    except Exception as e: return f"Error: {e}", None
+    except Exception as e:
+        return f"Error: {e}", None
 
 def upscale_image(image_path, factor):
     try:
         image = Image.open(image_path)
         new_size = (image.width * factor, image.height * factor)
         return image.resize(new_size, Image.LANCZOS)
-    except Exception as e: st.error(f"Error: {e}")
+    except Exception as e:
+        st.error(f"Error: {e}")
 
 def save_image(image, seed):
     try:
         path = DATA_PATH / f"image_{seed}.jpg"
         image.save(path, format="JPEG")
         return path
-    except Exception as e: st.error(f"Error: {e}")
+    except Exception as e:
+        st.error(f"Error: {e}")
 
 def save_prompt(prompt, seed):
     try:
         path = DATA_PATH / f"prompt_{seed}.txt"
-        with open(path, 'w') as f: f.write(prompt)
+        with open(path, 'w') as f:
+            f.write(prompt)
         return path
-    except Exception as e: st.error(f"Error: {e}")
+    except Exception as e:
+        st.error(f"Error: {e}")
 
 async def gen(prompt, model, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
     combined_prompt = f"{prompt} {await improve_prompt(prompt, language)}" if process_enhancer else prompt
@@ -64,7 +75,7 @@ async def gen(prompt, model, width, height, scales, steps, seed, upscale_factor,
     progress_bar.progress(50)
     if isinstance(image, str) and image.startswith("Error"): return [image, None]
     image_path = save_image(image, seed)
-    prompt_path = save_prompt(combined_prompt, seed)
+    prompt_path = save_prompt(combined_prompt, seed)  
     if process_upscale: image = upscale_image(image_path, upscale_factor)
     progress_bar.progress(100)
     return [str(image_path) if not process_upscale else str(save_image(image, seed)), str(prompt_path)]
@@ -75,13 +86,17 @@ async def improve_prompt(prompt, language):
     response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
     return response.get('generated_text', '').strip()[:500]
 
-def authenticate_user(username, password, credentials): 
+def authenticate_user(username, password, credentials):
     return username == credentials["username"] and password == credentials["password"]
 
 def login_form(credentials):
     username, password = st.text_input("Usuario"), st.text_input("Contraseña", type="password")
-    if st.button("Iniciar Sesión") and authenticate_user(username, password, credentials): 
-        st.session_state['authenticated'] = True
+    if st.button("Iniciar Sesión"):
+        if authenticate_user(username, password, credentials):
+            st.session_state['authenticated'] = True
+            st.session_state['username'] = username  
+        else:
+            st.error("Credenciales incorrectas")
 
 def main():
     st.set_page_config(layout="wide")
@@ -89,7 +104,8 @@ def main():
         st.title("Iniciar Sesión")
         login_form(credentials)
         return
-    
+
+    # Mantener la sesión activa
     st.sidebar.text_input("Descripción de la imagen", max_chars=900, key='prompt')
     process_enhancer = st.sidebar.checkbox("Mejorar Prompt", value=False)
     language = st.sidebar.selectbox("Idioma", ["en", "es"])
@@ -103,11 +119,14 @@ def main():
         result = run_async(gen(st.session_state['prompt'], basemodel, 512, 512, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language))
         if result[0] and isinstance(result[0], str):
             st.image(result[0], caption='Imagen Generada')
-            st.write(f"Prompt: {result[1]}")  
-            st.button("Borrar Imagen", on_click=lambda: st.session_state.clear())
-            st.button("Face Swap", on_click=lambda: st.write("Funcionalidad de Face Swap no implementada"))
+            st.write(f"Prompt: {result[1]}") 
+            if st.button("Borrar Imagen"):
+                os.remove(result[0]) 
+                st.success("Imagen borrada.")
+            if st.button("Face Swap"):
+                st.write("Funcionalidad de Face Swap no implementada")  
     
-    if 'authenticated' in st.session_state and st.session_state['authenticated']:
+    if st.sidebar.button("Cerrar Sesión"):
         st.session_state.clear()  
 
 if __name__ == "__main__":