Update app.py

app.py

@@ -1,80 +1,75 @@
 from pathlib import Path
 from PIL import Image
 import streamlit as st
-import insightface
-from insightface.app import FaceAnalysis
-from huggingface_hub import InferenceClient
 import os, random, numpy as np, yaml, time, logging
 from dataclasses import dataclass, field
 from typing import List
+from huggingface_hub import InferenceClient
+
 st.set_page_config(layout="wide")
+
 try:
-    with open("config.yaml", "r") as file: credentials = yaml.safe_load(file)
+    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": ""}
+
 @dataclass
 class AppConfig:
     MAX_SEED: int = 1000000  
     CLEANUP_DAYS: int = 7    
+
 MAX_SEED = AppConfig.MAX_SEED
 client = InferenceClient()
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN")
-@st.cache_resource
-def prepare_face_app():
-    app = FaceAnalysis(name='buffalo_l')
-    app.prepare(ctx_id=0, det_size=(640, 640))
-    swapper = insightface.model_zoo.get_model('onix.onnx')
-    return app, swapper
-def get_upscale_finegrain(prompt, img_path, upscale_factor):
-    try:
-        upscale_client = InferenceClient("fal/AuraSR-v2", hf_token=HF_TOKEN_UPSCALER)
-        result = upscale_client.predict(input_image=open(img_path, "rb").read(), prompt=prompt, upscale_factor=upscale_factor)
-        return result[1] if isinstance(result, list) and len(result) > 1 else None
-    except Exception as e:
-        st.error(f"Error al mejorar la imagen: {e}")
-        return None
+
 def cleanup_old_images(max_age_days=AppConfig.CLEANUP_DAYS):
     current_time = time.time()
     for image_file in DATA_PATH.glob("*.jpg"):
-        if current_time - image_file.stat().st_mtime > max_age_days * 86400: os.remove(image_file)
-def authenticate_user(username, password): return username == credentials["username"] and password == credentials["password"]
-def list_saved_images(): return sorted(DATA_PATH.glob("*.jpg"), key=lambda x: x.stat().st_mtime, reverse=True)
-def sort_faces(faces): return sorted(faces, key=lambda x: x.bbox[0])
-def get_face(faces, face_id):
-    if not faces or len(faces) < face_id: raise ValueError("Rostro no disponible.")
-    return faces[face_id - 1]
-def swap_faces(source_image, source_face_index, destination_image, destination_face_index):
-    faces = sort_faces(app.get(source_image))
-    source_face = get_face(faces, source_face_index)
-    res_faces = sort_faces(app.get(destination_image))
-    if destination_face_index > len(res_faces) or destination_face_index < 1: raise ValueError("Índice de rostro de destino no válido.")
-    res_face = get_face(res_faces, destination_face_index)
-    result = swapper.get(destination_image, res_face, source_face, paste_back=True)
-    return result
+        if current_time - image_file.stat().st_mtime > max_age_days * 86400: 
+            os.remove(image_file)
+
+def authenticate_user(username, password): 
+    return username == credentials["username"] and password == credentials["password"]
+
+def list_saved_images(): 
+    return sorted(DATA_PATH.glob("*.jpg"), key=lambda x: x.stat().st_mtime, reverse=True)
+
 def enhance_prompt(text):
     try:
-        enhanced = client.text_generation("With this text, generate a descriptive and photorealistic txt2img prompt in English in 200 characters maximum: " + text,
-            model="mistralai/Mixtral-8x7B-v0.1", max_length=200)
+        enhanced = client.text_generation(
+            "With this text, generate a descriptive and photorealistic txt2img prompt in English in 200 characters maximum: " + text,
+            model="mistralai/Mixtral-8x7B-v0.1", 
+            max_length=200
+        )
         return enhanced[:200]
-    except: return text[:200]
+    except: 
+        return text[:200]
+
 def save_prompt(image_name, prompt):
-    with open(DATA_PATH / "prompts.txt", "a") as f: f.write(f"{image_name}: {prompt}\n")
+    with open(DATA_PATH / "prompts.txt", "a") as f: 
+        f.write(f"{image_name}: {prompt}\n")
+
 def generate_variations(prompt, num_variants=8, use_enhanced=True):
-    instructions = ["Create a photorealistic description for a detailed txt2img prompt in English: ",
+    instructions = [
+        "Create a photorealistic description for a detailed txt2img prompt in English: ",
         "Write a creative, realistic, and detailed text-to-image prompt in English: ",
         "Generate a descriptive and true to life txt2img prompt in English: ",
         "Describe a photorealistic scene with detailed illumination for a txt2img prompt: ",
         "Give a realistic, elegant txt2img prompt in English, emphasizing photorealism: ",
         "Create a visually dynamic and hyperrealistic txt2img prompt in English: ",
         "Write a cinematic txt2img prompt in English with hyperrealistic elements: ",
-        "Make a lifelike txt2img prompt in English, focusing on photorealistic depth: "]
+        "Make a lifelike txt2img prompt in English, focusing on photorealistic depth: "
+    ]
     if use_enhanced:
         prompts = [enhance_prompt(f"{instructions[i % len(instructions)]}{prompt}") for i in range(num_variants)]
-    else: prompts = [prompt] * num_variants
+    else: 
+        prompts = [prompt] * num_variants
     return prompts
+
 def generate_image(prompt, width, height, seed, model_name):
     try:
         with st.spinner("Generando imagen..."):
@@ -85,12 +80,15 @@ def generate_image(prompt, width, height, seed, model_name):
     except Exception as e:
         st.error(f"Error al generar imagen: {e}")
         return None, seed, None
+
 def gen(prompts, width, height, model_name, num_variants=8):
     images = []
     seeds = []
     while len(seeds) < num_variants:
         seed = random.randint(0, MAX_SEED)
-        if seed not in seeds: seeds.append(seed)
+        if seed not in seeds: 
+            seeds.append(seed)
+    
     for i in range(num_variants):
         current_prompt = prompts[i] if len(prompts) > i else prompts[-1]
         with st.spinner(f"Generando imagen {i+1}/{num_variants}"):
@@ -99,37 +97,37 @@ def gen(prompts, width, height, model_name, num_variants=8):
                 image_path = DATA_PATH / f"generated_image_{used_seed}.jpg"
                 image.save(image_path)
                 save_prompt(f"generated_image_{used_seed}.jpg", enhanced_prompt)
-                images.append(str(image_path))
+                images.append((str(image_path), enhanced_prompt))
                 st.success(f"Imagen {i+1} generada")
     return images
+
 def display_gallery():
     st.header("Galería de Imágenes Guardadas")
     images = list_saved_images()
-    image_options = {str(image_file): f"Imagen {i+1}" for i, image_file in enumerate(images)}
-    selected_image = st.radio("Seleccionar imagen para FaceSwap", options=image_options.keys(), index=None)    
-    if selected_image:
-        st.session_state['generated_image_path'] = selected_image
-        prompt = get_prompt_for_image(Path(selected_image).name)
-        st.write(f"Prompt: {prompt[:100]}")
-    if selected_image and st.button("Aplicar FaceSwap"):
-        st.success("Imagen seleccionada para FaceSwap")
+    
     if images:
         cols = st.columns(4) 
         for i, image_file in enumerate(images):
             with cols[i % 4]:
                 st.image(str(image_file), use_column_width=True)
+                prompt = get_prompt_for_image(image_file.name)
+                st.caption(prompt[:100])
                 if st.button(f"Borrar", key=f"delete_{i}_{image_file}"):
                     if image_file.exists():
                         os.remove(image_file)
                         st.success("Imagen borrada")
                         st.rerun()
+
 def get_prompt_for_image(image_name):
     try:
         with open(DATA_PATH / "prompts.txt", "r") as f:
             for line in f:
-                if line.startswith(image_name): return line.split(": ", 1)[1].strip()
-    except FileNotFoundError: return "No hay prompt asociado"
+                if line.startswith(image_name): 
+                    return line.split(": ", 1)[1].strip()
+    except FileNotFoundError: 
+        return "No hay prompt asociado"
     return "No hay prompt asociado"
+
 def login_form():
     st.title("Iniciar Sesión")
     username = st.text_input("Usuario", value="admin")
@@ -138,7 +136,9 @@ def login_form():
         if authenticate_user(username, password):
             st.session_state['authenticated'] = True
             st.success("Autenticación exitosa.")
-        else: st.error("Credenciales incorrectas. Intenta de nuevo.")
+        else: 
+            st.error("Credenciales incorrectas. Intenta de nuevo.")
+
 def upload_image_to_gallery():
     uploaded_image = st.sidebar.file_uploader("Sube una imagen a la galería", type=["jpg", "jpeg", "png"])
     if uploaded_image:
@@ -147,27 +147,37 @@ def upload_image_to_gallery():
         image.save(image_path)
         save_prompt(uploaded_image.name, "uploaded by user")
         st.sidebar.success(f"Imagen subida: {image_path}")
-app, swapper = prepare_face_app()
+
 def main():
     if 'authenticated' not in st.session_state or not st.session_state['authenticated']: 
         login_form()
         return
-    st.title("Flux +Upscale +Prompt Enhancer +FaceSwap")
-    generated_image_path = st.session_state.get('generated_image_path')
+    
+    st.title("Flux +Upscale +Prompt Enhancer")
+    
     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"])
     model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"])
-    upscale_checkbox = st.sidebar.checkbox("Escalar imagen")
     prompt_enhance = st.sidebar.checkbox("Mejorar Prompt", True)
     num_variants = st.sidebar.slider("Número de imágenes", 1, 8, 8)
+    
     width, height = (720, 1280) if format_option == "9:16" else (1280, 720) if format_option == "16:9" else (1280, 1280)
+    
     if prompt:
         prompts = generate_variations(prompt, num_variants=num_variants, use_enhanced=prompt_enhance)
-        if st.sidebar.button("Generar Imágenes"): images = gen(prompts, width, height, model_option, num_variants)
-    if generated_image_path and upscale_checkbox:
-        upscale_factor = st.sidebar.slider("Factor de Escalado", 1, 4, 2)
-        improved_image = get_upscale_finegrain(prompt, generated_image_path, upscale_factor)
-        if improved_image: st.image(improved_image, caption="Imagen Escalada", use_column_width=True)
+        if st.sidebar.button("Generar Imágenes"):
+            generated_images = gen(prompts, width, height, model_option, num_variants)
+            
+            # Display generated images with their prompts
+            st.header("Imágenes Generadas")
+            cols = st.columns(4)
+            for i, (image_path, image_prompt) in enumerate(generated_images):
+                with cols[i % 4]:
+                    st.image(image_path, use_column_width=True)
+                    st.caption(image_prompt)
+    
     upload_image_to_gallery()
     display_gallery()
-if __name__ == "__main__": main()
+
+if __name__ == "__main__": 
+    main()