Update app.py

app.py

@@ -1,204 +1,102 @@
-import os
-import zipfile
-import yaml
-import numpy as np
-import random
+import os, zipfile, yaml, numpy as np, random, asyncio
 from pathlib import Path
 from PIL import Image
 import streamlit as st
 from huggingface_hub import AsyncInferenceClient
-import asyncio
 from moviepy.editor import ImageSequenceClip
 
-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": ""} 
+try: credentials = yaml.safe_load(open("config.yaml"))
+except: st.error("Error al cargar el archivo de configuración."), credentials = {"username": "", "password": ""}
 
-MAX_SEED = np.iinfo(np.int32).max
-client = AsyncInferenceClient()
-DATA_PATH = Path("./data")
+MAX_SEED, client, DATA_PATH, PREDEFINED_SEED = np.iinfo(np.int32).max, AsyncInferenceClient(), Path("./data"), random.randint(0, MAX_SEED)
 DATA_PATH.mkdir(exist_ok=True)
-PREDEFINED_SEED = random.randint(0, MAX_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=model_name
-        )
-        return image, seed
-    except Exception as e:
-        return f"Error al generar imagen: {e}", None
-
-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 prompt_file_path
-    except Exception as e:
-        st.error(f"Error al guardar el prompt: {e}")
-        return None
-
-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, model_name)
+
+async def generate_image(prompt, width, height, seed, model_name): 
+    try: return await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name), int(seed)
+    except Exception as e: return f"Error al generar imagen: {e}", None
+
+def save_prompt(prompt_text, seed): 
+    try: prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"; open(prompt_file_path, "w").write(prompt_text); return prompt_file_path
+    except Exception as e: st.error(f"Error al guardar el prompt: {e}")
+
+async def gen(prompt, width, height, model_name): 
+    seed, progress_bar = PREDEFINED_SEED, st.progress(0)
+    image, seed = await generate_image(prompt, width, height, seed, model_name)
     progress_bar.progress(100)
+    if isinstance(image, str) and image.startswith("Error"): return [image, None]
+    return [str(save_image(image, seed)), str(save_prompt(prompt, seed))]
+
+def save_image(image, seed): 
+    try: image_path = DATA_PATH / f"image_{seed}.jpg"; image.save(image_path, format="JPEG"); return image_path
+    except Exception as e: st.error(f"Error al guardar la imagen: {e}")
+
+def get_storage(): 
+    files = sorted([file for file in DATA_PATH.glob("*.jpg")], key=lambda x: x.stat().st_mtime, reverse=True)
+    return [str(file.resolve()) for file in files], f"Uso total: {sum([file.stat().st_size for file in files])/(1024.0 ** 3):.3f}GB"
+
+def get_prompts(): return {file.stem.replace("prompt_", ""): file for file in DATA_PATH.glob("*.txt")}
+
+def delete_all_images(): 
+    try: [os.remove(file) for file in DATA_PATH.glob("*.jpg") + DATA_PATH.glob("*.txt")], 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"; zipf = zipfile.ZipFile(zip_path, 'w')
+    [zipf.write(file, arcname=file.name) for file in DATA_PATH.glob("*.jpg")]
+    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")
 
-    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 image_path
-    except Exception as e:
-        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 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 create_video_from_images():
-    try:
-        image_files = sorted(DATA_PATH.glob("*.jpg"))
-        if not image_files:
-            st.error("No hay imágenes disponibles para crear un video.")
-            return
-        
-        image_sequence = [Image.open(image_file) for image_file in image_files]
-        frame_rate = 2
-        
-        clip = ImageSequenceClip([np.array(img) for img in image_sequence], fps=1)  
-        video_path = DATA_PATH / "output_video.mp4"
-        clip.write_videofile(str(video_path), codec="libx264")
-
-        return video_path  
-        
-    except Exception as e:
-        st.error(f"Error al generar el video: {e}")
-        return None
+def create_video_from_images(): 
+    try: image_files = sorted(DATA_PATH.glob("*.jpg"))
+    except: st.error("No hay imágenes disponibles para crear un video."), return
+    image_sequence = [Image.open(image_file) for image_file in image_files]
+    video_path = DATA_PATH / "output_video.mp4"
+    ImageSequenceClip([np.array(img) for img in image_sequence], fps=1).write_videofile(str(video_path), codec="libx264")
+    return video_path
 
 def main():
     st.set_page_config(layout="wide")
-    st.title("Generador de Imágenes")
+    st.title("Generador de Imágenes Flux")
     
-    if "authenticated" not in st.session_state:
-        st.session_state.authenticated = False
+    if "authenticated" not in st.session_state: st.session_state.authenticated = False
 
     if not st.session_state.authenticated:
-        st.subheader("Iniciar sesión")
-        username = st.text_input("Usuario")
-        password = st.text_input("Contraseña", type="password")
-
-        if st.button("Ingresar"):
-            if username == credentials["username"] and password == credentials["password"]:
-                st.session_state.authenticated = True
-                st.success("Inicio de sesión exitoso.")
-            else:
-                st.error("Usuario o contraseña incorrectos.")
+        username, password = st.text_input("Usuario"), st.text_input("Contraseña", type="password")
+        if st.button("Ingresar") and username == credentials["username"] and password == credentials["password"]: 
+            st.session_state.authenticated, st.success("Inicio de sesión exitoso.") = True, True
+        elif st.button("Ingresar"): st.error("Usuario o contraseña incorrectos.")
         return
 
-    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", ["enhanceaiteam/Flux-Uncensored-V2", "enhanceaiteam/Flux-uncensored"])
-
-        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, model_option))
-                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 st.sidebar.button("Borrar todas las imágenes"):
-            delete_all_images()
-
-        if st.sidebar.button("Descargar imágenes en .zip"):
-            download_images_as_zip()
-
-        if st.button("Generar video con las imágenes"):
-            video_path = create_video_from_images()
-            if video_path:
-                st.video(str(video_path), format="video/mp4")
-
-if __name__ == "__main__":
-    main()
+    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", ["enhanceaiteam/Flux-Uncensored-V2", "enhanceaiteam/Flux-uncensored"])
+    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, model_option))
+            image_paths, prompt_file = result[0], result[1]
+            if Path(image_paths).exists(): st.image(image_paths, caption="Imagen Generada")
+            if prompt_file and Path(prompt_file).exists(): st.write(f"Prompt utilizado: {Path(prompt_file).read_text()}")
+
+    files, usage = get_storage()
+    st.text(usage)
+    cols, prompts = st.columns(6), 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)
+            st.image(image, caption=f"Imagen {idx+1}")
+            st.write(f"Prompt: {Path(prompt_file).read_text() if prompt_file else 'No disponible'}")
+            if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"): 
+                os.remove(file)
+                if prompt_file: os.remove(prompt_file)
+                st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
+
+    if st.sidebar.button("Borrar todas las imágenes"): delete_all_images()
+    if st.sidebar.button("Descargar imágenes en .zip"): download_images_as_zip()
+    if st.button("Generar video con las imágenes"): 
+        video_path = create_video_from_images()
+        if video_path: st.video(str(video_path), format="video/mp4")
+
+if __name__ == "__main__": main()