Update app.py

app.py

@@ -1,210 +1,162 @@
-import os
-import numpy as np
-import random
+import os, random, asyncio, numpy as np
 from pathlib import Path
 from PIL import Image
+from insightface.app import FaceAnalysis
 import streamlit as st
 from huggingface_hub import InferenceClient, AsyncInferenceClient
 from gradio_client import Client, handle_file
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
 import yaml
-import cv2
-import dlib
-
-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": ""}
+import insightface 
 
 MAX_SEED = np.iinfo(np.int32).max
+DATA_PATH = Path("./data"); DATA_PATH.mkdir(exist_ok=True)
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
-client = AsyncInferenceClient()
-llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-DATA_PATH = Path("./data")
-DATA_PATH.mkdir(exist_ok=True)
-
-def run_async(func):
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    executor = ThreadPoolExecutor(max_workers=1)
-    result = loop.run_in_executor(executor, func)
-    return loop.run_until_complete(result)
-
-async def generate_image(combined_prompt, model, width, height, scales, steps, seed):
-    seed = int(seed) if seed != -1 else random.randint(0, MAX_SEED)
-    try:
-        image = await client.text_to_image(
-            prompt=combined_prompt, height=height, width=width, guidance_scale=scales,
-            num_inference_steps=steps, model=model
-        )
-        return image, seed
-    except Exception as e:
-        return f"Error al generar imagen: {e}", None
-
-def get_upscale_finegrain(prompt, img_path, upscale_factor):
-    try:
-        client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
-        result = 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:
-        return None
+client, llm_client = AsyncInferenceClient(), InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
-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)
-        return prompt_file_path
-    except Exception as e:
-        st.error(f"Error al guardar el prompt: {e}")
-        return None
+try: 
+    credentials = yaml.safe_load(open("config.yaml"))
+except Exception as e: 
+    st.error(f"Error al cargar config: {e}"); credentials = {"username": "", "password": ""}
 
-async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
-    combined_prompt = f"{prompt} {await improve_prompt(prompt, language) if process_enhancer else ''}".strip()
-    seed = int(seed) if seed != -1 else random.randint(0, MAX_SEED)
-    progress_bar = st.progress(0)
+def prepare_face_app():
+    app = FaceAnalysis(name='buffalo_l'); app.prepare(ctx_id=0, det_size=(640, 640))
+    return app, insightface.model_zoo.get_model('onix.onnx')
 
-    image, seed = await generate_image(combined_prompt, basemodel, width, height, scales, steps, seed)
-    progress_bar.progress(50)
+app, swapper = prepare_face_app()
 
-    if isinstance(image, str) and image.startswith("Error"):
-        progress_bar.empty()
-        return [image, None, combined_prompt]
+async def generate_image(prompt, model, w, h, scale, steps, seed):
+    seed = random.randint(0, MAX_SEED) if seed == -1 else seed
+    image = await client.text_to_image(prompt=prompt, height=h, width=w, guidance_scale=scale, num_inference_steps=steps, model=model)
+    return image, seed if not isinstance(image, str) else (None, None)
 
-    image_path = save_image(image, seed)
-    prompt_file_path = save_prompt(combined_prompt, seed)
+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
 
+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)
     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")
-            progress_bar.progress(100)
             image_path.unlink()
             return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
-    
-    progress_bar.progress(100)
     return [str(image_path), str(prompt_file_path)]
 
 async def improve_prompt(prompt, language):
-    instruction = (
-        "Con esta idea, describe en español un prompt detallado de txt2img en un máximo de 500 caracteres, "
-        "con iluminación, atmósfera, elementos cinematográficos y en su caso personajes..."
-        if language == "es" else
-        "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, "
-        "add illumination, atmosphere, cinematic elements, and characters if needed..."
-    )
-    formatted_prompt = f"{prompt}: {instruction}"
-    response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
-    improved_text = response.get('generated_text', '').strip() if 'generated_text' in response else response.strip()
-    return improved_text[:500] if len(improved_text) > 500 else improved_text
+    instruction = "With this idea, describe in English a detailed txt2img prompt in 500 characters at most..." if language == "en" else "Con esta idea, describe en español un prompt detallado de txt2img..."
+    response = await llm_client.text_generation(f"{prompt}: {instruction}", max_new_tokens=500)
+    return response.get('generated_text', '').strip()[:500]
 
 def save_image(image, seed):
-    try:
-        image_path = DATA_PATH / f"image_{seed}.jpg"
-        image.save(image_path, format="JPEG")
-        return image_path
-    except Exception as e:
-        st.error(f"Error al guardar la imagen: {e}")
-        return None
+    if image.mode == 'RGBA': image = image.convert('RGB')
+    image_path = DATA_PATH / f"image_{seed}.jpg"
+    image.save(image_path, format="JPEG")
+    return image_path
+
+def save_prompt(prompt_text, seed):
+    prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
+    open(prompt_file_path, "w").write(prompt_text)
+    return prompt_file_path
 
 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)
-    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}
+    total_size = sum([file.stat().st_size for file in files]) / (1024.0 ** 3)
+    return files, f"Uso total: {total_size:.3f} GB"
 
 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}")
+
+def delete_all_images():
+    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 not st.session_state['authenticated']:
+        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.")
+
+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())
+        st.sidebar.success(f"Imagen {uploaded_file.name} cargada correctamente.")
+        return image_path, save_prompt("#uploadedbyuser", image_path.stem)
+    return None
+
+def gallery():
+    files, usage = get_storage()
+    st.sidebar.write(f"{usage}")
+    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]:
+            st.image(str(file))
+            try:
+                prompt_file_path = DATA_PATH / f"prompt_{file.stem.split('_')[-1]}.txt"
+                st.write(f"Prompt: {open(prompt_file_path).read()}")
+            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)
+
+def face_swap(image_path, source_image_path):
     try:
-        if Path(image_path).exists():
-            Path(image_path).unlink()
-            st.success(f"Imagen {image_path} borrada.")
-        else:
-            st.error("El archivo de imagen no existe.")
-    except Exception as e:
-        st.error(f"Error al borrar la imagen: {e}")
-
-def swap_faces(image_path):
-    try:
-        image = cv2.imread(str(image_path))
-        detector = dlib.get_frontal_face_detector()
-        predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
-        faces = detector(image)
-        
-        if len(faces) != 2:
-            st.error("Se necesitan exactamente dos caras para realizar el intercambio.")
-            return None
-
-        landmarks1 = predictor(image, faces[0])
-        landmarks2 = predictor(image, faces[1])
-        points1 = np.array([[p.x, p.y] for p in landmarks1.parts()])
-        points2 = np.array([[p.x, p.y] for p in landmarks2.parts()])
-
-        mask1 = np.zeros(image.shape[:2], dtype=np.uint8)
-        cv2.fillConvexPoly(mask1, cv2.convexHull(points1), 255)
-        mask2 = np.zeros(image.shape[:2], dtype=np.uint8)
-        cv2.fillConvexPoly(mask2, cv2.convexHull(points2), 255)
-        face1 = cv2.bitwise_and(image, image, mask=mask1)
-        face2 = cv2.bitwise_and(image, image, mask=mask2)
-        image[mask1 == 255] = face2[mask1 == 255]
-        image[mask2 == 255] = face1[mask2 == 255]
-        swapped_image_path = DATA_PATH / f"swapped_image_{Path(image_path).stem}.jpg"
-        cv2.imwrite(str(swapped_image_path), image)
-        return str(swapped_image_path)
-    except Exception as e:
-        st.error(f"Error en el face swap: {e}")
-        return None
-
-def main():
+        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])
+        swapped_img_path = DATA_PATH / f"swapped_{Path(image_path).stem}.jpg"
+        swapped_img.save(swapped_img_path, format="JPEG")
+        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)
+
+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
     prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=900)
-    process_enhancer = st.sidebar.checkbox("Mejorar Prompt", value=False)
-    language = st.sidebar.selectbox("Idioma", ["en", "es"])
-    basemodel = st.sidebar.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-DEV", "black-forest-labs/FLUX.1-schnell"])
-    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
-    process_upscale = st.sidebar.checkbox("Procesar Escalador", value=False)
-    upscale_factor = st.sidebar.selectbox("Factor de Escala", [2, 4, 8], index=0)
-    scales = st.sidebar.slider("Escalado", 1, 20, 10)
-    steps = st.sidebar.slider("Pasos", 1, 100, 20)
-    seed = st.sidebar.number_input("Semilla", value=-1)
-
-    width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
+    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)
+    upload_image()
 
+    image_path, prompt_file_path = None, None  
     if st.sidebar.button("Generar Imagen"):
-        with st.spinner("Mejorando y generando imagen..."):
-            result = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language))
-            image_paths, prompt_file = result[0], result[1]
-
-        st.write(f"Image paths: {image_paths}")
-
-        if image_paths and Path(image_paths).exists():
-            st.image(image_paths, caption="Imagen generada", use_column_width=True)
-        if prompt_file and Path(prompt_file).exists():
-            with open(prompt_file, "r") as file:
-                st.text_area("Prompt utilizado", file.read(), height=150)
-
-    st.sidebar.header("Galería de Imágenes")
-    image_storage, usage = get_storage()
-    st.sidebar.write(usage)
-    for img_path in image_storage:
-        st.sidebar.image(img_path, width=100)
+        with st.spinner("Generando..."):
+            image_path, prompt_file_path = await gen(prompt, basemodel, w, h, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language)
+            if image_path:
+                st.image(image_path, caption="Imagen Generada")
+                st.write(f"Prompt: {open(prompt_file_path).read()}")
     
-    if st.sidebar.button("Borrar Imagen"):
-        delete_image(image_paths)
-
-    if st.sidebar.button("Intercambiar Caras"):
-        if image_paths:
-            swapped_path = swap_faces(image_paths)
-            if swapped_path:
-                st.image(swapped_path, caption="Imagen con caras intercambiadas", use_column_width=True)
+    if image_path:
+        st.success("Imagen generada y almacenada.")
+    
+    gallery()
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())