Update app.py

app.py

@@ -1,4 +1,7 @@
-import os, random, asyncio, numpy as np
+import os
+import random
+import asyncio
+import numpy as np
 from pathlib import Path
 from PIL import Image
 from insightface.app import FaceAnalysis
@@ -6,20 +9,23 @@ import streamlit as st
 from huggingface_hub import InferenceClient, AsyncInferenceClient
 from gradio_client import Client, handle_file
 import yaml
-import insightface 
+import insightface
 
 MAX_SEED = np.iinfo(np.int32).max
-DATA_PATH = Path("./data"); DATA_PATH.mkdir(exist_ok=True)
+DATA_PATH = Path("./data")
+DATA_PATH.mkdir(exist_ok=True)
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
 client, llm_client = AsyncInferenceClient(), InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 try: 
     credentials = yaml.safe_load(open("config.yaml"))
 except Exception as e: 
-    st.error(f"Error al cargar config: {e}"); credentials = {"username": "", "password": ""}
+    st.error(f"Error al cargar config: {e}")
+    credentials = {"username": "", "password": ""}
 
 def prepare_face_app():
-    app = FaceAnalysis(name='buffalo_l'); app.prepare(ctx_id=0, det_size=(640, 640))
+    app = FaceAnalysis(name='buffalo_l')
+    app.prepare(ctx_id=0, det_size=(640, 640))
     return app, insightface.model_zoo.get_model('onix.onnx')
 
 app, swapper = prepare_face_app()
@@ -33,18 +39,20 @@ def get_upscale_finegrain(prompt, img_path, upscale_factor):
     try:
         result = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER).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: return None
+    except Exception:
+        return None
 
 async def gen(prompt, basemodel, w, h, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
     combined_prompt = f"{prompt} {await improve_prompt(prompt, language)}" if process_enhancer else prompt
     image, seed = await generate_image(combined_prompt, basemodel, w, h, scales, steps, seed)
     if image is None: return ["Error al generar imagen", None, combined_prompt]
-    image_path = save_image(image, seed); prompt_file_path = save_prompt(combined_prompt, seed)
+    image_path = save_image(image, seed)
+    prompt_file_path = save_prompt(combined_prompt, seed)
     if process_upscale:
         upscale_image_path = get_upscale_finegrain(combined_prompt, image_path, upscale_factor)
         if upscale_image_path:
             Image.open(upscale_image_path).save(DATA_PATH / f"upscale_image_{seed}.jpg", format="JPEG")
-            image_path.unlink()
+            image_path.unlink()  # Borra la imagen original
             return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
     return [str(image_path), str(prompt_file_path)]
 
@@ -71,29 +79,38 @@ def get_storage():
 
 def delete_image(image_path):
     try: 
-        Path(image_path).unlink(); st.success(f"Imagen {image_path} borrada.")
-    except Exception as e: st.error(f"Error al borrar imagen: {e}")
+        Path(image_path).unlink()
+        st.success(f"Imagen {image_path} borrada.")
+    except Exception as e: 
+        st.error(f"Error al borrar imagen: {e}")
 
 def delete_all_images():
-    for file in DATA_PATH.glob("*.jpg"): file.unlink(); st.success("Todas las imágenes han sido borradas.")
+    for file in DATA_PATH.glob("*.jpg"): 
+        file.unlink()
+    st.success("Todas las imágenes han sido borradas.")
 
 def authenticate_user(username, password, credentials):
     return username == credentials["username"] and password == credentials["password"]
 
 def login_form(credentials):
-    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']:
-        username = st.text_input("Usuario"); password = st.text_input("Contraseña", type='password')
+        username = st.text_input("Usuario")
+        password = st.text_input("Contraseña", type='password')
         if st.button("Iniciar Sesión"):
             if authenticate_user(username, password, credentials):
-                st.session_state['authenticated'] = True; st.success("Inicio de sesión exitoso.")
-            else: st.error("Credenciales incorrectas.")
+                st.session_state['authenticated'] = True
+                st.success("Inicio de sesión exitoso.")
+            else: 
+                st.error("Credenciales incorrectas.")
 
 def upload_image():
     uploaded_file = st.sidebar.file_uploader("Sube una imagen", type=["png", "jpg", "jpeg"])
     if uploaded_file:
         image_path = DATA_PATH / uploaded_file.name
-        with open(image_path, "wb") as f: f.write(uploaded_file.getbuffer())
+        with open(image_path, "wb") as f: 
+            f.write(uploaded_file.getbuffer())
         st.sidebar.success(f"Imagen {uploaded_file.name} cargada correctamente.")
         return image_path, save_prompt("#uploadedbyuser", image_path.stem)
     return None
@@ -101,7 +118,8 @@ def upload_image():
 def gallery():
     files, usage = get_storage()
     st.sidebar.write(f"{usage}")
-    if st.sidebar.button("Borrar Todas las Imágenes"): delete_all_images()
+    if st.sidebar.button("Borrar Todas las Imágenes"): 
+        delete_all_images()
     cols = st.columns(6)
     for idx, file in enumerate(files):
         with cols[idx % 6]:
@@ -112,37 +130,51 @@ def gallery():
             except FileNotFoundError:
                 st.write("Prompt no encontrado.")
             st.button(f"Borrar Imagen {file.name}", on_click=delete_image, args=(file,))
-            if st.button(f"Swap Face en {file.name}"): upload_source_and_swap(file)
+            if st.button(f"Swap Face en {file.name}"): 
+                upload_source_and_swap(file)
 
-def face_swap(image_path, source_image_path):
+def upload_source_and_swap(image_path):
+    source_image = st.file_uploader("Sube la imagen source para face swap", type=["png", "jpg", "jpeg"])
+    if source_image:
+        source_image_path = DATA_PATH / source_image.name
+        with open(source_image_path, "wb") as f: 
+            f.write(source_image.getbuffer())
+        st.success(f"Imagen source {source_image.name} cargada correctamente.")
+        dest_face_index = st.slider("Selecciona la posición de la cara destino", 1, 4, 1)
+        swapped_image_path = face_swap(image_path, source_image_path, dest_face_index)
+        if swapped_image_path: 
+            st.image(str(swapped_image_path), caption="Imagen con Face Swap", use_column_width=True)
+
+def face_swap(image_path, source_image_path, dest_face_index):
     try:
         img_dest, img_src = Image.open(image_path), Image.open(source_image_path)
         faces = app.get(img_src)
-        if not faces: st.error("No se encontraron caras en la imagen source."); return None
-        swapped_img = swapper.get(img_dest, faces[0])
+        if not faces: 
+            st.error("No se encontraron caras en la imagen source.")
+            return None
+        swapped_img = swapper.get(img_dest, faces[0], dest_index=dest_face_index - 1)  # Convertir a índice base 0
         swapped_img_path = DATA_PATH / f"swapped_{Path(image_path).stem}.jpg"
         swapped_img.save(swapped_img_path, format="JPEG")
+        Path(image_path).unlink() 
         return swapped_img_path
-    except Exception as e: st.error(f"Error en face swap: {e}"); return None
-
-def upload_source_and_swap(image_path):
-    source_image = st.file_uploader("Sube la imagen source para face swap", type=["png", "jpg", "jpeg"])
-    if source_image:
-        source_image_path = DATA_PATH / source_image.name
-        with open(source_image_path, "wb") as f: f.write(source_image.getbuffer())
-        st.success(f"Imagen source {source_image.name} cargada correctamente.")
-        swapped_image_path = face_swap(image_path, source_image_path)
-        if swapped_image_path: st.image(str(swapped_image_path), caption="Imagen con Face Swap", use_column_width=True)
+    except Exception as e: 
+        st.error(f"Error en face swap: {e}")
+        return None
 
 async def main():
     st.set_page_config(layout="wide")
     login_form(credentials)
-    if not st.session_state['authenticated']: st.warning("Por favor, inicia sesión para acceder a la aplicación."); return
+    if not st.session_state['authenticated']:
+        st.warning("Por favor, inicia sesión para acceder a la aplicación.")
+        return
+
     prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=900)
     process_enhancer, language = st.sidebar.checkbox("Mejorar Prompt", value=False), st.sidebar.selectbox("Idioma", ["en", "es"])
     basemodel, format_option, process_upscale = st.sidebar.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-DEV", "black-forest-labs/FLUX.1-schnell"]), st.sidebar.selectbox("Formato", ["9:16", "16:9"]), st.sidebar.checkbox("Procesar Escalador", value=False)
     upscale_factor, scales, steps, seed = st.sidebar.selectbox("Factor de Escala", [2, 4, 8], index=0), st.sidebar.slider("Escalado", 1, 20, 10), st.sidebar.slider("Pasos", 1, 100, 20), st.sidebar.number_input("Semilla", value=-1)
-    w, h = (1080, 1920) if format_option == "9:16" else (1920, 1080)
+    destination_face_position = st.sidebar.slider("Posición de la Cara de Destino", 1, 4, 1)
+    
+    w, h = (720, 1280) if format_option == "9:16" else (1280, 720)
     upload_image()
 
     image_path, prompt_file_path = None, None  
@@ -159,4 +191,4 @@ async def main():
     gallery()
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())