Update app.py

app.py

@@ -1,10 +1,10 @@
 import os
+import zipfile
 import numpy as np
 import random
 from pathlib import Path
 from PIL import Image
 import streamlit as st
-import yaml
 from huggingface_hub import AsyncInferenceClient
 import asyncio
 
@@ -13,23 +13,15 @@ client = AsyncInferenceClient()
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 
-# Cargar credenciales desde archivo config.yaml
-try:
-    with open("config.yaml", "r") as file:
-        credentials = yaml.safe_load(file)
-except Exception as e:
-    st.error(f"Error al cargar el archivo de configuración: {e}")
-    credentials = {"username": "", "password": ""}  # Valores predeterminados
-
 PREDEFINED_SEED = random.randint(0, MAX_SEED)
 
-async def generate_image(prompt, width, height, seed):
+async def generate_image(prompt, width, height, seed, model_name):
     try:
         if seed == -1:
             seed = PREDEFINED_SEED
         seed = int(seed)
         image = await client.text_to_image(
-            prompt=prompt, height=height, width=width, model="enhanceaiteam/Flux-uncensored"
+            prompt=prompt, height=height, width=width, model=model_name
         )
         return image, seed
     except Exception as e:
@@ -45,11 +37,11 @@ def save_prompt(prompt_text, seed):
         st.error(f"Error al guardar el prompt: {e}")
         return None
 
-async def gen(prompt, width, height):
+async def gen(prompt, width, height, model_name):
     combined_prompt = prompt  
     seed = PREDEFINED_SEED
     progress_bar = st.progress(0)
-    image, seed = await generate_image(combined_prompt, width, height, seed)
+    image, seed = await generate_image(combined_prompt, width, height, seed, model_name)
     progress_bar.progress(100)
 
     if isinstance(image, str) and image.startswith("Error"):
@@ -80,31 +72,40 @@ def get_prompts():
     prompt_files = [file for file in DATA_PATH.glob("*.txt") if file.is_file()]
     return {file.stem.replace("prompt_", ""): file for file in prompt_files}
 
+def delete_all_images():
+    try:
+        files = [file for file in DATA_PATH.glob("*.jpg")]
+        prompts = [file for file in DATA_PATH.glob("*.txt")]
+        for file in files + prompts:
+            os.remove(file)
+        st.success("Todas las imágenes y prompts han sido borrados.")
+    except Exception as e:
+        st.error(f"Error al borrar archivos: {e}")
+
+def download_images_as_zip():
+    zip_path = DATA_PATH / "images.zip"
+    with zipfile.ZipFile(zip_path, 'w') as zipf:
+        for file in DATA_PATH.glob("*.jpg"):
+            zipf.write(file, arcname=file.name)
+    with open(zip_path, "rb") as zip_file:
+        st.download_button(label="Descargar imágenes en .zip", data=zip_file, file_name="images.zip", mime="application/zip")
+
 def main():
     st.set_page_config(layout="wide")
-    st.title("Flux Uncensored")
-    
-    if "logged_in" not in st.session_state:
-        st.session_state.logged_in = False
-
-    if not st.session_state.logged_in:
-        username = st.text_input("Usuario")
-        password = st.text_input("Contraseña", type="password")
-        if st.button("Iniciar Sesión"):
-            if username == credentials["username"] and password == credentials["password"]:
-                st.session_state.logged_in = True
-                st.success("Inicio de sesión exitoso.")
-            else:
-                st.error("Usuario o contraseña incorrectos.")
-    else:
+    st.title("Generador de Imágenes")
+    st.warning("Este espacio contiene contenido que no es adecuado para todas las audiencias. Se recomienda discreción.")
+    agree = st.checkbox("Soy mayor de 18 años y entiendo que el contenido puede no ser apropiado.")
+
+    if agree:
         prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
         format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
+        model_option = st.sidebar.selectbox("Modelo", ["modelo_actual", "enhanceaiteam/Flux-Uncensored-V2"])
         
         width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
 
         if st.sidebar.button("Generar Imagen"):
             with st.spinner("Generando imagen..."):
-                result = asyncio.run(gen(prompt, width, height))
+                result = asyncio.run(gen(prompt, width, height, model_option))
                 image_paths = result[0]
                 prompt_file = result[1]
 
@@ -143,5 +144,13 @@ def main():
                     except Exception as e:
                         st.error(f"Error al borrar la imagen o prompt: {e}")
 
+        if st.button("Borrar todas las imágenes"):
+            delete_all_images()
+
+        download_images_as_zip()
+
+    else:
+        st.error("No puedes acceder a este espacio hasta que confirmes que eres mayor de edad.")
+
 if __name__ == "__main__": 
     main()