Update app.py

app.py

@@ -1,47 +1,175 @@
-import base64
-import io
-from huggingface_hub import InferenceClient
-import streamlit as st
+from pathlib import Path
 from PIL import Image
+import streamlit as st
+from huggingface_hub import InferenceClient, AsyncInferenceClient
+import asyncio
+import os
+import random
+import numpy as np
+import 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": ""}
+
+MAX_SEED = np.iinfo(np.int32).max
+client = AsyncInferenceClient()
+llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+DATA_PATH = Path("./data")
+DATA_PATH.mkdir(exist_ok=True)
 
-# Configuración
-hf_token = st.secrets["HF_TOKEN"]
-client = InferenceClient(repo_id="black-forest-labs/FLUX.1-schnell", token=hf_token)
+def authenticate_user(username, password):
+    return username == credentials["username"] and password == credentials["password"]
+
+async def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=True):
+    images = []
+    try:
+        for idx, prompt in enumerate(prompts):
+            seed = random.randint(0, MAX_SEED)
+            image, seed = await generate_image(prompt, width, height, seed, model_name)
+            image_path = save_image(image, f"generated_image_{seed}.jpg", prompt)
+            if image_path:
+                st.success(f"Imagen {idx + 1} generada")
+                images.append(str(image_path))
+    except Exception as e:
+        st.error(f"Error al generar imágenes: {e}")
+    return images
 
-def get_image_result(prompt, init_image=None):
-    if init_image:
-        img_buffer = io.BytesIO()
-        init_image.save(img_buffer, format="PNG")
-        img_bytes = img_buffer.getvalue()
-        encoded_image = base64.b64encode(img_bytes).decode('utf-8')
+def list_saved_images():
+    return sorted(DATA_PATH.glob("*.jpg"), key=os.path.getmtime, reverse=True)
+
+def display_gallery():
+    st.header("Galería de Imágenes Guardadas")
+    images = list_saved_images()
+    if images:
+        cols = st.columns(8)
+        for i, image_file in enumerate(images):
+            with cols[i % 8]:
+                st.image(str(image_file), caption=image_file.name, use_column_width=True)
+                prompt = get_prompt_for_image(image_file.name)
+                st.write(prompt[:300])
+
+                if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"):
+                    os.remove(image_file)
+                    st.success("Imagen borrada")
+                    display_gallery()
     else:
-        encoded_image = None
+        st.info("No hay imágenes guardadas.")
+
+def save_prompt(prompt):
+    with open(DATA_PATH / "prompts.txt", "a") as f:
+        f.write(prompt + "\n")
 
+def run_async(func, *args):
+    return asyncio.run(func(*args))
+
+async def improve_prompt(prompt):
     try:
-        output = client.img2img(prompt=prompt, image=encoded_image, strength=0.75)
-        img_data = base64.b64decode(output['generated_image'])
-        return Image.open(io.BytesIO(img_data))
+        instructions = [
+            "With my idea create a vibrant description for a detailed txt2img prompt, 300 characters max.",
+            "With my idea write a creative and detailed text-to-image prompt in English, 300 characters max.",
+            "With my idea generate a descriptive and visual txt2img prompt in English, 300 characters max.",
+            "With my idea describe a photorealistic with illumination txt2img prompt in English, 300 characters max.",
+            "With my idea give a realistic and elegant txt2img prompt in English, 300 characters max.",
+            "With my idea conform a visually dynamic and surreal txt2img prompt in English, 300 characters max.",
+            "With my idea realize an artistic and cinematic txt2img prompt in English, 300 characters max.",
+            "With my idea make a narrative and immersive txt2img prompt in English, 300 characters max."
+        ]
+        instruction = random.choice(instructions)
+        formatted_prompt = f"{prompt}: {instruction}"
+        response = llm_client.text_generation(formatted_prompt, max_new_tokens=100)
+        return response['generated_text'][:100] if 'generated_text' in response else response.strip()
     except Exception as e:
-        st.error(f"Error en la generación de la imagen: {str(e)}")
+        return f"Error mejorando el prompt: {e}"
+
+def save_image(image, file_name, prompt=None):
+    image_path = DATA_PATH / file_name
+    if image_path.exists():
+        st.warning(f"La imagen '{file_name}' ya existe en la galería. No se guardó.")
         return None
+    else:
+        image.save(image_path, format="JPEG")
+        if prompt:
+            save_prompt(f"{file_name}: {prompt}")
+        return image_path
 
-# Interfaz de Streamlit
-st.title("Generación de Imágenes img2img con Flux")
-st.sidebar.header("Opciones de generación")
-prompt = st.sidebar.text_input("Escribe tu prompt:", value="Una escena futurista en una ciudad verde")
-uploaded_file = st.sidebar.file_uploader("Sube una imagen base (opcional)", type=["jpg", "png"])
+async def generate_image(prompt, width, height, seed, model_name):
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
+    image = await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
+    return image, seed
 
-if st.sidebar.button("Generar Imagen"):
-    if uploaded_file:
-        init_image = Image.open(uploaded_file).convert("RGB")
-        st.image(init_image, caption="Imagen Base", use_column_width=True)
+def get_prompt_for_image(image_name):
+    prompts = {}
+    try:
+        with open(DATA_PATH / "prompts.txt", "r") as f:
+            for line in f:
+                if line.startswith(image_name):
+                    prompts[image_name] = line.split(": ", 1)[1].strip()
+    except FileNotFoundError:
+        return "No hay prompt asociado."
+    return prompts.get(image_name, "No hay prompt asociado.")
+
+def login_form():
+    st.title("Iniciar Sesión")
+    username = st.text_input("Usuario", value="admin")
+    password = st.text_input("Contraseña", value="flux3x", type="password")
+    if st.button("Iniciar Sesión"):
+        if authenticate_user(username, password):
+            st.success("Autenticación exitosa.")
+            st.session_state['authenticated'] = True
+        else:
+            st.error("Credenciales incorrectas. Intenta de nuevo.")
+
+async def generate_variations(prompt, num_variants, use_enhanced):
+    prompts = set() 
+    while len(prompts) < num_variants:
+        if use_enhanced:
+            enhanced_prompt = await improve_prompt(prompt)
+            prompts.add(enhanced_prompt)
+        else:
+            prompts.add(prompt)
+    return list(prompts)
+
+async def main():
+    st.set_page_config(layout="wide")
+
+    if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
+        login_form()
+        return
+
+    st.title("Flux + Multiple Images")
+    prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500)
+    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9", "1:1"]) 
+    prompt_checkbox = st.sidebar.checkbox("Prompt Enhancer")
+    model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"])
+    width, height = (360, 640) if format_option == "9:16" else (640, 360) if format_option == "16:9" else (640, 640)
+
+    if prompt_checkbox:
+        num_variants = st.sidebar.slider("Número de imágenes a generar", 1, 8, 1)
+    else:
+        num_variants = 1
+
+    if prompt_checkbox:
+        with st.spinner("Generando prompts mejorados..."):
+            prompts = await generate_variations(prompt, num_variants, True)
     else:
-        init_image = None
-
-    with st.spinner('Generando imagen...'):
-        generated_image = get_image_result(prompt, init_image)
-        if generated_image:
-            st.image(generated_image, caption="Imagen Generada", use_column_width=True)
-            img_bytes = io.BytesIO()
-            generated_image.save(img_bytes, format="PNG")
-            st.download_button(label="Descargar imagen", data=img_bytes.getvalue(), file_name="generated_image.png", mime="image/png")
+        prompts = [prompt] 
+
+    if st.sidebar.button("Generar Imágenes"):
+        with st.spinner("Generando imágenes..."):
+            try:
+                results = await gen(prompts, width, height, model_option, num_variants, prompt_checkbox)
+                st.session_state['generated_image_paths'] = results
+                for result in results:
+                    st.image(result, caption="Imagen Generada")
+            except Exception as e:
+                st.error(f"Error al generar las imágenes: {str(e)}")
+
+    display_gallery()
+
+if __name__ == "__main__":
+    asyncio.run(main())