Update app.py

app.py

@@ -1,204 +1,186 @@
-import os
-import zipfile
-import yaml
-import numpy as np
-import random
 from pathlib import Path
 from PIL import Image
 import streamlit as st
-from huggingface_hub import AsyncInferenceClient
-import asyncio
-from moviepy.editor import ImageSequenceClip
+import insightface
+from insightface.app import FaceAnalysis
+from huggingface_hub import InferenceClient
+import os
+import random
+import numpy as np
+import yaml
 
-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": ""}
+with open("config.yaml", "r") as file:
+    credentials = yaml.safe_load(file)
 
 MAX_SEED = np.iinfo(np.int32).max
-client = AsyncInferenceClient()
+client = InferenceClient()
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 PREDEFINED_SEED = random.randint(0, MAX_SEED)
-
-async def generate_image(prompt, width, height, seed, model_name):
-    try:
-        if seed == -1:
-            seed = PREDEFINED_SEED
-        seed = int(seed)
-        image = await client.text_to_image(
-            prompt=prompt, height=height, width=width, model=model_name
-        )
-        return image, seed
-    except Exception as e:
-        return f"Error al generar imagen: {e}", None
-
-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
-
-async def gen(prompt, width, height, model_name):
-    combined_prompt = prompt  
-    seed = PREDEFINED_SEED
-    progress_bar = st.progress(0)
-    image, seed = await generate_image(combined_prompt, width, height, seed, model_name)
-    progress_bar.progress(100)
-
-    if isinstance(image, str) and image.startswith("Error"):
-        progress_bar.empty()
-        return [image, None]
-
-    image_path = save_image(image, seed)
-    prompt_file_path = save_prompt(combined_prompt, seed)
-
-    return [str(image_path), str(prompt_file_path)]
-
-def save_image(image, seed):
+HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN")
+
+def get_upscale_finegrain(prompt, img_path, upscale_factor):
+    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
+
+def authenticate_user(username, password):
+    return username == credentials["username"] and password == credentials["password"]
+
+def list_saved_images():
+    image_files = sorted(DATA_PATH.glob("*.jpg"))
+    return image_files
+
+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):
+    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))
+    res_face = get_face(res_faces, destination_face_index)
+    result = swapper.get(destination_image, res_face, source_face, paste_back=True)
+    return result
+
+def generate_image(prompt, width, height, seed, model_name):
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
+    image = client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
+    return image, seed
+
+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[:100])
+                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 image_file.exists():
+                        os.remove(image_file)
+                        st.success("Imagen borrada")
+                    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 generate_variations(prompt, num_variants, use_enhanced):
+    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"
+    ]
+    prompts = set()
+    while len(prompts) < num_variants:
+        instruction = random.choice(instructions)
+        enhanced_prompt = f"{instruction}: {prompt}"
+        prompts.add(enhanced_prompt)
+    return list(prompts)
+
+def gen(prompts, width, height, model_name, num_variants=1):
+    images = []
+    seeds = random.sample(range(MAX_SEED), num_variants)
+    for idx, (prompt, seed) in enumerate(zip(prompts[:num_variants], seeds)):
+        image, _ = generate_image(prompt, width, height, seed, model_name)
+        image_path = save_image(image, f"generated_image_{seed}.jpg")
+        if image_path:
+            save_prompt(f"generated_image_{seed}.jpg: {prompt}")
+            st.success(f"Imagen {idx + 1} generada")
+            images.append(str(image_path))
+    return images
+
+def get_prompt_for_image(image_name):
+    prompts = {}
     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
-
-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}
-
-def delete_all_images():
-    try:
-        files = [file for file in DATA_PATH.glob("*.jpg")]
-        prompts = [file for file in DATA_PATH.glob("*.txt")]
-        for file in files + prompts:
-            os.remove(file)
-        st.success("Todas las imágenes y prompts han sido borrados.")
-    except Exception as e:
-        st.error(f"Error al borrar archivos: {e}")
-
-def download_images_as_zip():
-    zip_path = DATA_PATH / "images.zip"
-    with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for file in DATA_PATH.glob("*.jpg"):
-            zipf.write(file, arcname=file.name)
-    with open(zip_path, "rb") as zip_file:
-        st.download_button(label="Descargar imágenes en .zip", data=zip_file, file_name="images.zip", mime="application/zip")
-
-def create_video_from_images():
-    try:
-        image_files = sorted(DATA_PATH.glob("*.jpg"))
-        if not image_files:
-            st.error("No hay imágenes disponibles para crear un video.")
-            return
-        
-        image_sequence = [Image.open(image_file) for image_file in image_files]
-        frame_rate = 2
-        
-        clip = ImageSequenceClip([np.array(img) for img in image_sequence], fps=1)  
-        video_path = DATA_PATH / "output_video.mp4"
-        clip.write_videofile(str(video_path), codec="libx264")
-
-        return video_path  
-        
-    except Exception as e:
-        st.error(f"Error al generar el video: {e}")
-        return None
+        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 save_image(image, filename):
+    image_path = DATA_PATH / filename
+    if isinstance(image, bytes):
+        with open(image_path, "wb") as f:
+            f.write(image)
+    else:
+        image.save(image_path)
+    return image_path
+
+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(f"{uploaded_image.name}: uploaded by user")
+            st.sidebar.success(f"Imagen subida: {image_path}")
 
 def main():
     st.set_page_config(layout="wide")
-    
-    if "authenticated" not in st.session_state:
-        st.session_state.authenticated = False
-
-    if not st.session_state.authenticated:
-        st.subheader("Iniciar sesión")
-        username = st.text_input("Usuario", value="admin")
-        password = st.text_input("Contraseña", value="flux3x", type="password") 
-
-        if st.button("Ingresar"):
-            if username == credentials["username"] and password == credentials["password"]:
-                st.session_state.authenticated = True
-                st.success("Inicio de sesión exitoso.")
-            else:
-                st.error("Usuario o contraseña incorrectos.")
+    if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
+        login_form()
         return
-
-    st.warning("Este espacio contiene contenido que no es adecuado para todas las audiencias. Se recomienda discreción.")
-    agree = st.checkbox("Soy mayor de 18 años y entiendo que el contenido puede no ser apropiado.")
-    st.title("Flux +Uncensored")
-    
-    if agree:
-        prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
-        format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
-        model_option = st.sidebar.selectbox("Modelo", ["enhanceaiteam/Flux-Uncensored-V2", "enhanceaiteam/Flux-uncensored"])
-
-        width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
-
-        if st.sidebar.button("Generar Imagen"):
-            with st.spinner("Generando imagen..."):
-                result = asyncio.run(gen(prompt, width, height, model_option))
-                image_paths = result[0]
-                prompt_file = result[1]
-
-            if image_paths:
-                if Path(image_paths).exists():
-                    st.image(image_paths, caption="Imagen Generada")
-                else:
-                    st.error("El archivo de imagen no existe.")
-
-                if prompt_file and Path(prompt_file).exists():
-                    prompt_text = Path(prompt_file).read_text()
-                    st.write(f"Prompt utilizado: {prompt_text}")
-                else:
-                    st.write("El archivo del prompt no está disponible.")
-
-        files, usage = get_storage()
-        st.text(usage)
-        cols = st.columns(6)
-        prompts = get_prompts()
-
-        for idx, file in enumerate(files):
-            with cols[idx % 6]:
-                image = Image.open(file)
-                prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
-                prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
-
-                st.image(image, caption=f"Imagen {idx+1}")
-                st.write(f"Prompt: {prompt_text}")
-
-                if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
-                    try:
-                        os.remove(file)
-                        if prompt_file:
-                            os.remove(prompt_file)
-                        st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
-                    except Exception as e:
-                        st.error(f"Error al borrar la imagen o prompt: {e}")
-
-        if st.sidebar.button("Borrar todas las imágenes"):
-            delete_all_images()
-
-        if st.sidebar.button("Descargar imágenes en .zip"):
-            download_images_as_zip()
-
-        if st.button("Generar video con las imágenes"):
-            video_path = create_video_from_images()
-            if video_path:
-                st.video(str(video_path), format="video/mp4")
+    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", "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)
+    upload_image_to_gallery()
+    display_gallery()
 
 if __name__ == "__main__":
-    main()
+    main()