Update app.py

app.py

@@ -1,79 +1,86 @@
-import os
-import random
 from pathlib import Path
-from PIL import Image, ImageOps
+from PIL import Image
 import streamlit as st
+import insightface
+from insightface.app import FaceAnalysis
 from huggingface_hub import InferenceClient, AsyncInferenceClient
-import face_recognition
+import asyncio
+import os
+import random
+import numpy as np
+import yaml
 
-TEMP_PATH = Path("./temp")
-TEMP_PATH.mkdir(parents=True, exist_ok=True)
-st.set_page_config(layout="wide")
+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)
+PREDEFINED_SEED = random.randint(0, MAX_SEED)
+HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
 
-def save_prompt(description):
-    prompt_file = TEMP_PATH / "prompts.txt"
-    with open(prompt_file, "a") as file:
-        file.write(f"{description}\n")
-
-def save_image(image, file_name, prompt=None):
-    image_path = TEMP_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
-    image.save(image_path, format="JPEG")
-    if prompt:
-        save_prompt(f"{file_name}: {prompt}")
-    return image_path
+if not HF_TOKEN_UPSCALER:
+    st.warning("HF_TOKEN_UPSCALER no está configurado. Algunas funcionalidades pueden no funcionar.")
 
-def generate_thumbnail(image_file, size=(150, 150)):
+def get_upscale_finegrain(prompt, img_path, upscale_factor):
     try:
-        with Image.open(image_file) as img:
-            img.thumbnail(size, Image.LANCZOS) 
-            thumbnail_path = TEMP_PATH / f"{Path(image_file).stem}_thumbnail.jpg"
-            img.save(thumbnail_path, format="JPEG")
-        return thumbnail_path
+        upscale_client = InferenceClient("fal/AuraSR-v2", hf_token=HF_TOKEN_UPSCALER)
+        result = upscale_client.predict(input_image=handle_file(img_path), 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 generar thumbnail: {e}")
+        st.error(f"Error al mejorar la imagen: {e}")
         return None
 
-def show_image_gallery(images):
-    cols = st.columns(8)
-    for i, image_file in enumerate(images):
-        with cols[i % 8]:
-            thumbnail_image = Image.open(image_file)
-            st.image(thumbnail_image, caption=f"Imagen {i + 1}", use_column_width=True)
-            if st.button(f"Eliminar Imagen {i + 1}", key=f"delete_{i + 1}"):
-                os.remove(image_file)
-                st.warning(f"Imagen {i + 1} eliminada")
-
-async def improve_prompt(prompt):
-    try:
-        instructions = [
-            "With this words, create a photorealistic description for a detailed txt2img prompt in English in 200 characters maximum",
-            "With this idea, write a creative, realistic, and detailed text-to-image prompt in English in 200 characters maximum",
-            "With this text, generate a descriptive and True to life txt2img prompt in English in 200 characters maximum",
-            "With my idea, describe a photorealistic scene with detailed illumination for a txt2img prompt in English in 200 characters maximum",
-            "With this concept, give a realistic, elegant txt2img prompt in English, emphasizing photorealism in 200 characters maximum",
-            "With this perspective, conform a visually dynamic and hyperrealistic txt2img prompt in English in 200 characters maximum",
-            "With this inspiration, realize a cinematic txt2img prompt in English with hyperrealistic elements in 200 characters maximum",
-            "With my idea, make a lifelike and txt2img prompt in English, focusing on photorealistic depth in 200 characters maximum"
-        ]
-
-        instruction = random.choice(instructions)
-        formatted_prompt = f"{prompt}: {instruction}"
-        response = llm_client.text_generation(formatted_prompt, max_new_tokens=256)
-        return response['generated_text'][:200] if 'generated_text' in response else response.strip()
-    except Exception as e:
-        return f"Error mejorando el prompt: {e}"
+def authenticate_user(username, password):
+    return username == credentials["username"] and password == credentials["password"]
+
+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
+
+app, swapper = prepare_face_app()
+
+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
+
+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
 
-async def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=True, seed=-1):
+async def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=True):
     images = []
     try:
         for idx, prompt in enumerate(prompts):
-            image, seed_used = await generate_image(prompt, width, height, model_name, seed)
-            image_path = save_image(image, f"generated_image_{seed_used}.jpg", prompt)
+            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")
             if image_path:
                 st.success(f"Imagen {idx + 1} generada")
                 images.append(str(image_path))
@@ -81,19 +88,87 @@ async def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=T
         st.error(f"Error al generar imágenes: {e}")
     return images
 
-async def generate_image(prompt, width, height, model_name, seed):
-    if seed == -1:
-        seed = random.randint(0, 2147483647)
-    image = await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name, seed=seed)
-    return image, seed
-
 def list_saved_images():
-    return sorted(TEMP_PATH.glob("*.jpg"), key=os.path.getmtime, reverse=True)
+    return list(DATA_PATH.glob("*.jpg"))
+
+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"FaceSwap", key=f"select_{i}_{image_file.name}"):
+                    st.session_state['generated_image_path'] = str(image_file)
+                    st.success("Imagen seleccionada")
+
+                if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"):
+                    if os.path.exists(image_file):
+                        os.remove(image_file)
+                        st.success("Imagen borrada")
+                        display_gallery()
+                    else:
+                        st.warning("La imagen no existe.")
+    else:
+        st.info("No hay imágenes guardadas.")
+
+def save_prompt(prompt):
+    with open(DATA_PATH / "prompts.txt", "a") as f:
+        f.write(prompt + "\n")
+        st.success("Prompt guardado.")
+
+def run_async(func, *args):
+    return asyncio.run(func(*args))
+
+async def improve_prompt(prompt):
+    try:
+        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:
+        return f"Error mejorando el prompt: {e}"
+
+def generate_variations(prompt, num_variants, use_enhanced):
+    prompts = set() 
+    while len(prompts) < num_variants:
+        if use_enhanced:
+            enhanced_prompt = improve_prompt(prompt)
+            prompts.add(enhanced_prompt)
+        else:
+            prompts.add(prompt)
+    return list(prompts)
+
+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")
+    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.")
@@ -101,92 +176,79 @@ def login_form():
         else:
             st.error("Credenciales incorrectas. Intenta de nuevo.")
 
-async def generate_variations(prompt, num_variants, use_enhanced, style):
-    prompts = set()
-    while len(prompts) < num_variants:
-        enhanced_prompt = await improve_prompt(f"{prompt}, estilo: {style}") if use_enhanced else f"{prompt}, estilo: {style}"
-        prompts.add(enhanced_prompt)
-    return list(prompts)
-
-def authenticate_user(username, password):
-    return username == "admin" and password == "flux3x"
-
-def swap_faces(source_image_path, target_image_path, output_path):
+def save_image(image, filename):
     try:
-        source_image = face_recognition.load_image_file(source_image_path)
-        target_image = face_recognition.load_image_file(target_image_path)
-
-        source_face_locations = face_recognition.face_locations(source_image)
-        target_face_locations = face_recognition.face_locations(target_image)
-
-        if not source_face_locations or not target_face_locations:
-            raise ValueError("No se detectaron rostros en una de las imágenes.")
-
-        source_face_encodings = face_recognition.face_encodings(source_image, source_face_locations)
-        target_face_encodings = face_recognition.face_encodings(target_image, target_face_locations)
-
-        # En este ejemplo, solo se realiza un intercambio básico del primer rostro encontrado
-        for target_location, target_encoding in zip(target_face_locations, target_face_encodings):
-            match = face_recognition.compare_faces(source_face_encodings, target_encoding, tolerance=0.6)
-            if any(match):
-                top, right, bottom, left = target_location
-                target_image[top:bottom, left:right] = source_image[top:bottom, left:right]
-
-        swapped_image = Image.fromarray(target_image)
-        swapped_image.save(output_path)
-        return output_path
-
+        image_path = DATA_PATH / filename
+        image.save(image_path)
+        return image_path
     except Exception as e:
-        st.error(f"Error al intercambiar rostros: {e}")
+        st.error(f"Error al guardar la imagen: {e}")
         return None
 
+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:
+        image = Image.open(uploaded_image)
+        image_path = save_image(image, f"{uploaded_image.name}")
+        if image_path:
+            save_prompt("uploaded by user")
+            st.sidebar.success(f"Imagen subida: {image_path}")
+
 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 +Upscale +Prompt Enhancer +FaceSwap")
+    generated_image_path = st.session_state.get('generated_image_path')
     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"])
-    width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
-    st.sidebar.subheader("Configuración de Imagen")
-    model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "prashanth970/flux-lora-uncensored"])
-    style_type = st.sidebar.selectbox("Estilo", ["Realismo", "Hiperrealismo", "Photorealismo"])
-    st.sidebar.subheader("Parámetros de Generación")
-    randomize_seed = st.sidebar.checkbox("Randomize Seed", value=True)
-    seed = -1 if randomize_seed else st.sidebar.number_input("Seed", value=42, min_value=0, max_value=2147483647)
-    prompt_checkbox = st.sidebar.checkbox("Mejorar Prompt", value=True)
-    num_variants = st.sidebar.slider("Número de imágenes a generar", 1, 8, 1)
-
+    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"])
+    prompt_checkbox = st.sidebar.checkbox("Mejorar Prompt")
+    upscale_checkbox = st.sidebar.checkbox("Escalar imagen")
+    width, height = (720, 1280) if format_option == "9:16" else (1280, 720) if format_option == "16:9" else (1280, 1280)
+    num_variants = st.sidebar.slider("Número de imágenes a generar", 1, 8, 1) if prompt_checkbox else 1
+    
     if prompt_checkbox:
-        prompts = await generate_variations(f"{prompt}, estilo: {style_type}", num_variants, prompt_checkbox, style_type)
+        with st.spinner("Generando prompts mejorados..."):
+            prompts = generate_variations(prompt, num_variants, True)
     else:
-        prompts = [f"{prompt}, estilo: {style_type}"]
+        prompts = [prompt] 
 
-    if st.sidebar.button("Generar Imágenes"):
-        try:
-            results = await gen(prompts, width, height, model_option, num_variants, prompt_checkbox, seed)
-            for result in results:
-                st.image(result, caption="Imagen Generada", use_column_width=True)
-        except Exception as e:
-            st.error(f"Error al generar las imágenes: {str(e)}")
+    upload_image_to_gallery()
 
-    images = list_saved_images()
-    if images:
-        show_image_gallery(images)
-
-        st.sidebar.subheader("Herramientas Adicionales")
-        source_image = st.sidebar.file_uploader("Cargar Imagen Fuente para FaceSwap", type=["jpg", "jpeg", "png"])
-        target_image_path = st.sidebar.selectbox("Seleccionar Imagen Objetivo", images)
-
-        if source_image and target_image_path:
-            output_path = TEMP_PATH / "face_swap_result.jpg"
-            if st.sidebar.button("Intercambiar Rostros"):
-                result = swap_faces(source_image, target_image_path, output_path)
-                if result:
-                    st.image(result, caption="Resultado de FaceSwap", use_column_width=True)
-    else:
-        st.info("No hay imágenes disponibles en la galería.")
+    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)}")
+
+    if generated_image_path:
+        if upscale_checkbox:
+            with st.spinner("Escalando imagen..."):
+                try:
+                    upscale_image_path = get_upscale_finegrain("Upscale", generated_image_path, 2)
+                    if upscale_image_path:
+                        st.image(upscale_image_path, caption="Imagen Escalada")
+                except Exception as e:
+                    st.error(f"Error al escalar la imagen: {str(e)}")
+
+        st.header("Intercambio de Rostros")
+        source_image_file = st.file_uploader("Imagen de Origen", type=["jpg", "jpeg", "png"])
+
+        if source_image_file is not None:
+            try:
+                source_image = Image.open(source_image_file)
+            except Exception as e:
+                st.error(f"Error al cargar la imagen de origen: {str(e)}")
+                source_image = None
+        else:
+            source_image = Image.open("face.jpg")
 
-if __name__ == "__main__":
-    import asyncio
-    asyncio.run(main())
+        source_face_index