Update app.py

app.py

@@ -1,201 +1,244 @@
-import os
-import random
-import asyncio
-import numpy as np
 from pathlib import Path
 from PIL import Image
-from insightface.app import FaceAnalysis
 import streamlit as st
+import insightface
+from insightface.app import FaceAnalysis
 from huggingface_hub import InferenceClient, AsyncInferenceClient
-from gradio_client import Client, handle_file
+import asyncio
+import os
+import random
+import numpy as np
 import yaml
-import insightface
+
+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")
-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": ""}
+def get_upscale_finegrain(prompt, img_path, upscale_factor):
+    try:
+        upscale_client = Client("finegrain/finegrain-image-enhancer", 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:
+        return None
+
+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))
-    return app, insightface.model_zoo.get_model('onix.onnx')
+    swapper = insightface.model_zoo.get_model('onix.onnx')
+    return app, swapper
 
 app, swapper = prepare_face_app()
 
-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)
+def sort_faces(faces):
+    return sorted(faces, key=lambda x: x.bbox[0])
 
-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
+def get_face(faces, face_id):
+    if not faces or len(faces) < face_id:
+        raise ValueError("Rostro no disponible.")
+    return faces[face_id - 1]
 
-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")
-            image_path.unlink() 
-            return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
-    return [str(image_path), str(prompt_file_path)]
-
-async def improve_prompt(prompt, language):
+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 = PREDEFINED_SEED
+    image = await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
+    return image, seed
+
+async def gen(prompt, width, height, model_name):
+    seed = PREDEFINED_SEED
+    image, seed = await generate_image(prompt, width, height, seed, model_name)
+    image_path = save_image(image, f"generated_image_{seed}.jpg", prompt)
+    return str(image_path)
+
+def list_saved_images():
+    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"Usar", 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}"):
+                    os.remove(image_file)
+                    st.success("Imagen borrada")
+                    display_gallery()
+    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:
-        instruction_en = "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters if need it..."
-        instruction_es = "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..."
-        instruction = instruction_en if language == "en" else instruction_es
+        instruction = ("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['generated_text'].strip() if 'generated_text' in response else response.strip()
-        return improved_text[:500] if len(improved_text) > 500 else improved_text
+        return response['generated_text'][:500] if 'generated_text' in response else response.strip()
     except Exception as e:
         return f"Error mejorando el prompt: {e}"
 
-def save_image(image, seed):
-    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()]
-    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 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):
+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
+
+def get_prompt_for_image(image_name):
+    prompts = {}
     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], 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
-
-async def main():
+        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.")
+
+def upload_image_to_gallery():
+    uploaded_image = st.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.success(f"Imagen subida: {image_path}")
+
+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.")
+
+    if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
+        login_form()
         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)
-    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()
+    st.title("Generador Flux")
+    generated_image_path = st.session_state.get('generated_image_path')
+    st.header("Generador de Imágenes")
+    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"])
+    model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"])
+    prompt_checkbox = st.sidebar.checkbox("Prompt Enhancer")
+    upscale_checkbox = st.sidebar.checkbox("Escalar imagen")
+    width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
+    upload_image_to_gallery()
+
+    if prompt_checkbox:
+        with st.spinner("Mejorando el prompt..."):
+            try:
+                improved_prompt = run_async(improve_prompt, prompt)
+            except Exception as e:
+                st.error(f"Error al mejorar el prompt: {str(e)}")
+                improved_prompt = prompt
+    else:
+        improved_prompt = prompt
 
-    image_path, prompt_file_path = None, None  
     if st.sidebar.button("Generar Imagen"):
-        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 image_path:
-        st.success("Imagen generada y almacenada.")
-    
-    gallery()
+        with st.spinner("Generando imagen..."):
+            try:
+                result = run_async(gen, improved_prompt, width, height, model_option)  # Usar el improved_prompt
+                st.session_state['generated_image_path'] = result
+                st.image(result, caption="Imagen Generada")
+            except Exception as e:
+                st.error(f"Error al generar la imagen: {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")
+
+        source_face_index = st.number_input('Posición del Rostro', min_value=1, value=1, key="source_face_index")
+        destination_face_index = st.number_input('Posición del Rostro de Destino', min_value=1, value=1, key="destination_face_index")
+
+        if st.button("Intercambiar Rostros"):
+            try:
+                destination_image = Image.open(generated_image_path)
+                result_image = swap_faces(np.array(source_image), source_face_index, np.array(destination_image), destination_face_index)
+                swapped_image = Image.fromarray(result_image)
+                
+                swapped_image_path = save_image(swapped_image, f"swapped_image_{PREDEFINED_SEED}.jpg")
+                if swapped_image_path:
+                    st.image(swapped_image, caption="Intercambio de Rostro")
+                    os.remove(generated_image_path)
+                else:
+                    st.warning("La imagen intercambiada ya existe en la galería.")
+            except Exception as e:
+                st.error(f"Ocurrió un error al intercambiar rostros: {str(e)}")
+
+    display_gallery()
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()