Update app.py

app.py

@@ -1,57 +1,30 @@
-import os
-import numpy as np
-import random
-from pathlib import Path
-from PIL import Image
-import streamlit as st
-from huggingface_hub import InferenceClient, AsyncInferenceClient
-import asyncio
+import os; import numpy as np; import random; from pathlib import Path; from PIL import Image; import streamlit as st; from huggingface_hub import AsyncInferenceClient; import asyncio; import yaml
 
-MAX_SEED = np.iinfo(np.int32).max
-client = AsyncInferenceClient()
-DATA_PATH = Path("./data")
-DATA_PATH.mkdir(exist_ok=True)
+with open("config.yaml") as f: config = yaml.safe_load(f); username = config["credentials"]["username"]; password = config["credentials"]["password"]
 
-PREDEFINED_SEED = random.randint(0, MAX_SEED)
+client = AsyncInferenceClient(); DATA_PATH = Path("./data"); DATA_PATH.mkdir(exist_ok=True); MAX_SEED = np.iinfo(np.int32).max; 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
+        seed = int(seed) if seed != -1 else PREDEFINED_SEED
+        return await client.text_to_image(prompt=prompt, height=height, width=width, model="enhanceaiteam/Flux-uncensored"), 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)
+        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
+        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)
+    seed = PREDEFINED_SEED; progress_bar = st.progress(0)
+    image, seed = 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(combined_prompt, seed)
-
-    return [str(image_path), str(prompt_file_path)]
+    if isinstance(image, str) and image.startswith("Error"): progress_bar.empty(); return [image, None]
+    return [save_image(image, seed), save_prompt(prompt, seed)]
 
 def save_image(image, seed):
     try:
@@ -59,72 +32,37 @@ def save_image(image, seed):
         image.save(image_path, format="JPEG")
         return image_path
     except Exception as e:
-        st.error(f"Error al guardar la imagen: {e}")
-        return None
+        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])
+    files = sorted(DATA_PATH.glob("*.jpg"), 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}
+    return {file.stem.replace("prompt_", ""): file for file in DATA_PATH.glob("*.txt") if file.is_file()}
 
 def main():
-    st.set_page_config(layout="wide")
-    st.title("Generador de Imágenes FLUX")
+    st.set_page_config(layout="wide"); st.title("Generador de Imágenes")
     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
+    width, height = (720, 1280) if st.sidebar.selectbox("Formato", ["9:16", "16:9"]) == "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()
+            if (image_paths := result[0]) and Path(image_paths).exists(): st.image(image_paths, caption="Imagen Generada")
+            else: st.error("El archivo de imagen no existe.")
+            if (prompt_file := result[1]) and Path(prompt_file).exists():
+                st.write(f"Prompt utilizado: {Path(prompt_file).read_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}")
-            
+            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}"):
-                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}")
+                try: os.remove(file); os.remove(prompt_file) if prompt_file else None; 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()
+if __name__ == "__main__": main()