Update app.py

app.py

@@ -4,33 +4,55 @@ import random
 from pathlib import Path
 from PIL import Image
 import streamlit as st
-import streamlit_authenticator as stauth
-from huggingface_hub import InferenceClient
-
-st.set_page_config(layout="wide")
-
-credentials = {
-    "usernames": {
-        "admin": {
-            "name": "admin",
-            "password": stauth.Hasher(["flux33"]).generate()[0]
-        }
-    }
-}
-
-authenticator = stauth.Authenticate(credentials, "imagen_gen_auth", "flux3", 30)
-client = InferenceClient()  
+from huggingface_hub import InferenceClient, AsyncInferenceClient
+import asyncio
+
+MAX_SEED = np.iinfo(np.int32).max
+client = AsyncInferenceClient()
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 
-def generate_image(prompt, width, height):
+PREDEFINED_SEED = random.randint(0, MAX_SEED)
+
+async def generate_image(prompt, width, height, seed):
+    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"
+        )
+        return image, seed
+    except Exception as e:
+        return f"Error al generar imagen: {e}", None
+
+def save_prompt(prompt_text, seed):
     try:
-        image = client.text_to_image(prompt=prompt, height=height, width=width, model="enhanceaiteam/Flux-uncensored")
-        return image
+        prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
+        with open(prompt_file_path, "w") as prompt_file:
+            prompt_file.write(prompt_text)
+        return prompt_file_path
     except Exception as e:
-        st.error(f"Error al generar imagen: {e}")
+        st.error(f"Error al guardar el prompt: {e}")
         return None
 
+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"
@@ -40,26 +62,69 @@ def save_image(image, seed):
         st.error(f"Error al guardar la imagen: {e}")
         return None
 
+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"
+
+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():
-    name, authentication_status, username = authenticator.login(
-        "Login", "main", fields={"username": "Username", "password": "Password", "login": "Login"})
-
-    if authentication_status:
-        st.title("Generador de Imágenes con FLUX")
-        
-        prompt = st.text_input("Introduce el prompt para generar la imagen")
-        if st.button("Generar Imagen"):
-            if prompt:
-                imagen_generada = generate_image(prompt, 512, 512) 
-                if imagen_generada:
-                    st.image(imagen_generada, caption="Imagen Generada", use_column_width=True)
+    st.set_page_config(layout="wide")
+    st.title("Generador de Imágenes FLUX")
+    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
+    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
+
+    if format_option == "9:16":
+        width = 720
+        height = 1280
+    else:
+        width = 1280
+        height = 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("Por favor, introduce un prompt.")
-    
-    elif authentication_status == False:
-        st.error('Nombre de usuario/contraseña incorrectos')
-    elif authentication_status == None:
-        st.warning('Por favor, introduce tu nombre de usuario y contraseña')
+                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()