Update app.py

app.py

@@ -4,135 +4,144 @@ import random
 from pathlib import Path
 from PIL import Image
 import streamlit as st
+import yaml
 from huggingface_hub import AsyncInferenceClient
 import asyncio
-import yaml
 
-# Configuración de cliente de inferencia
+MAX_SEED = np.iinfo(np.int32).max
 client = AsyncInferenceClient()
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 
-# Cargar credenciales desde archivo YAML
-with open("config.yaml", "r") as file:
-    credentials = yaml.safe_load(file)
+# 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
 
-# Variables de estado
-if 'logged_in' not in st.session_state:
-    st.session_state.logged_in = False
-if 'prompt' not in st.session_state:
-    st.session_state.prompt = ""
+PREDEFINED_SEED = random.randint(0, MAX_SEED)
 
-# Función de generación de imágenes
 async def generate_image(prompt, width, height, seed):
     try:
         if seed == -1:
-            seed = random.randint(0, np.iinfo(np.int32).max)
+            seed = PREDEFINED_SEED
+        seed = int(seed)
         image = await client.text_to_image(
             prompt=prompt, height=height, width=width, model="enhanceaiteam/Flux-uncensored"
         )
-        return image
+        return image, seed
     except Exception as e:
-        return f"Error al generar imagen: {e}"
+        return f"Error al generar imagen: {e}", None
 
-# Función para guardar prompts
 def save_prompt(prompt_text, seed):
     try:
         prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
         with open(prompt_file_path, "w") as prompt_file:
             prompt_file.write(prompt_text)
-        return str(prompt_file_path)
+        return prompt_file_path
     except Exception as e:
         st.error(f"Error al guardar el prompt: {e}")
         return None
 
-# Función para guardar imágenes
+async def gen(prompt, width, height):
+    combined_prompt = prompt  
+    seed = PREDEFINED_SEED
+    progress_bar = st.progress(0)
+    image, seed = await generate_image(combined_prompt, width, height, seed)
+    progress_bar.progress(100)
+
+    if isinstance(image, str) and image.startswith("Error"):
+        progress_bar.empty()
+        return [image, None]
+
+    image_path = save_image(image, seed)
+    prompt_file_path = save_prompt(combined_prompt, seed)
+
+    return [str(image_path), str(prompt_file_path)]
+
 def save_image(image, seed):
     try:
         image_path = DATA_PATH / f"image_{seed}.jpg"
         image.save(image_path, format="JPEG")
-        return str(image_path)
+        return image_path
     except Exception as e:
         st.error(f"Error al guardar la imagen: {e}")
         return None
 
-# Función para obtener archivos almacenados
 def get_storage():
     files = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
     files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
     usage = sum([file.stat().st_size for file in files])
     return [str(file.resolve()) for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
 
-# Función principal
-async def gen(prompt, width, height):
-    seed = random.randint(0, np.iinfo(np.int32).max)
-    progress_bar = st.progress(0)
-    image = await generate_image(prompt, width, height, seed)
-    progress_bar.progress(100)
-
-    if isinstance(image, str) and image.startswith("Error"):
-        progress_bar.empty()
-        return [image, None]
-
-    image_path = save_image(image, seed)
-    prompt_file_path = save_prompt(prompt, seed)
-
-    return [image_path, prompt_file_path]
+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 main():
     st.set_page_config(layout="wide")
     st.title("Generador de Imágenes FLUX")
+    
+    if "logged_in" not in st.session_state:
+        st.session_state.logged_in = False
 
-    # Inicio de sesión
     if not st.session_state.logged_in:
-        username = st.sidebar.text_input("Usuario")
-        password = st.sidebar.text_input("Contraseña", type="password")
-
-        if st.sidebar.button("Iniciar Sesión"):
+        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("Credenciales incorrectas.")
-        return
-
-    # Opciones después de iniciar sesión
-    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
-    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
-
-    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))
-            image_paths = result[0]
-            prompt_file = result[1]
-
-        if image_paths and Path(image_paths).exists():
-            st.image(image_paths, caption="Imagen Generada")
-            if prompt_file and Path(prompt_file).exists():
-                prompt_text = Path(prompt_file).read_text()
-                st.write(f"Prompt utilizado: {prompt_text}")
-        else:
-            st.error("El archivo de imagen no existe.")
-
-    files, usage = get_storage()
-    st.text(usage)
-    cols = st.columns(6)
-
-    for idx, file in enumerate(files):
-        with cols[idx % 6]:
-            image = Image.open(file)
-            prompt_text = Path(DATA_PATH / f"prompt_{file.stem.replace('image_', '')}.txt").read_text() if Path(DATA_PATH / f"prompt_{file.stem.replace('image_', '')}.txt").exists() else "No disponible"
-            st.image(image, caption=f"Imagen {idx+1}")
-            st.write(f"Prompt: {prompt_text}")
-
-            if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
-                try:
-                    os.remove(file)
-                    os.remove(DATA_PATH / f"prompt_{file.stem.replace('image_', '')}.txt")
-                    st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
-                except Exception as e:
-                    st.error(f"Error al borrar la imagen o prompt: {e}")
-
-if __name__ == "__main__":
+                st.error("Usuario o contraseña incorrectos.")
+    else:
+        prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
+        format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
+        
+        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))
+                image_paths = result[0]
+                prompt_file = result[1]
+
+            if image_paths:
+                if Path(image_paths).exists():
+                    st.image(image_paths, caption="Imagen Generada")
+                else:
+                    st.error("El archivo de imagen no existe.")
+
+                if prompt_file and Path(prompt_file).exists():
+                    prompt_text = Path(prompt_file).read_text()
+                    st.write(f"Prompt utilizado: {prompt_text}")
+                else:
+                    st.write("El archivo del prompt no está disponible.")
+
+        files, usage = get_storage()
+        st.text(usage)
+        cols = st.columns(6)
+        prompts = get_prompts()
+
+        for idx, file in enumerate(files):
+            with cols[idx % 6]:
+                image = Image.open(file)
+                prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
+                prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
+                
+                st.image(image, caption=f"Imagen {idx+1}")
+                st.write(f"Prompt: {prompt_text}")
+                
+                if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
+                    try:
+                        os.remove(file)  
+                        if prompt_file:
+                            os.remove(prompt_file)
+                        st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
+                    except Exception as e:
+                        st.error(f"Error al borrar la imagen o prompt: {e}")
+
+if __name__ == "__main__": 
     main()