Update app.py

app.py

@@ -1,22 +1,20 @@
 from pathlib import Path
 from PIL import Image
 import streamlit as st
-import insightface
-from insightface.app import FaceAnalysis
-from huggingface_hub import InferenceClient
-import os
-import random
-import numpy as np
-import yaml
-import time
-import logging
-from dataclasses import dataclass, field
+import os, random, numpy as np, yaml, time
+from dataclasses import dataclass
 from typing import List
+from huggingface_hub import InferenceClient
 
 st.set_page_config(layout="wide")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+if not HF_TOKEN:
+    st.error("Error en el token! 'HF_TOKEN'.")
+    st.stop() 
 
 try:
-    with open("config.yaml", "r") as file:
+    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}")
@@ -28,149 +26,93 @@ class AppConfig:
     CLEANUP_DAYS: int = 7    
 
 MAX_SEED = AppConfig.MAX_SEED
-client = InferenceClient()
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
-HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN")
 
-@st.cache_resource
-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
+def get_inference_client():
+    return InferenceClient(token=HF_TOKEN)
 
-def get_upscale_finegrain(prompt, img_path, upscale_factor):
-    try:
-        upscale_client = InferenceClient("fal/AuraSR-v2", hf_token=HF_TOKEN_UPSCALER)
-        result = upscale_client.predict(input_image=open(img_path, "rb").read(), 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 mejorar la imagen: {e}")
-        return None
-
-def cleanup_old_images(max_age_days=AppConfig.CLEANUP_DAYS):
-    current_time = time.time()
-    for image_file in DATA_PATH.glob("*.jpg"):
-        if current_time - image_file.stat().st_mtime > max_age_days * 86400:
-            os.remove(image_file)
+client = get_inference_client()
 
-def authenticate_user(username, password):
+def authenticate_user(username, password): 
     return username == credentials["username"] and password == credentials["password"]
 
-def list_saved_images():
+def list_saved_images(): 
     return sorted(DATA_PATH.glob("*.jpg"), key=lambda x: x.stat().st_mtime, reverse=True)
 
-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
-
-def generate_image(prompt, width, height, seed, model_name):
+def enhance_prompt(text, client=client):
+    if not client:
+        return text[:200]
     try:
-        with st.spinner("Generando imagen..."):
-            seed = int(seed) if seed != -1 else random.randint(0, AppConfig.MAX_SEED)
-            image = client.text_to_image(
-                prompt=prompt,
-                height=height,
-                width=width,
-                model=model_name,
-                seed=seed
-            )
-            return image, seed
+        enhanced = client.text_generation(
+            "Generate a photorealistic, detailed txt2img prompt: " + text,
+            model="mistralai/Mixtral-8x7B-Instruct-v0.1",)
+        return enhanced[:200]
+        
     except Exception as e:
-        st.error(f"Error al generar imagen: {e}")
-        return None, 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")
-                        display_gallery()
-
-def save_prompt(image_name, prompt):
-    with open(DATA_PATH / "prompts.txt", "a") as f:
-        f.write(f"{image_name}: {prompt[:100]}\n")
+        st.warning(f"Prompt enhancement error: {e}")
+        return text[:200]
 
-def enhance_prompt(text):
-    try:
-        enhanced = client.text_generation(
-            "With this text, generate a descriptive and photorealistic txt2img prompt in English in 200 characters maximum: " + text,
-            model="mistralai/Mixtral-8x7B-v0.1",
-            max_length=200
-        )
-        return enhanced[:100]
-    except:
-        return text[:100]
+def save_prompt(image_name, enhanced_prompt):
+    with open(DATA_PATH / "prompts.txt", "a") as f: 
+        f.write(f"{image_name}: {enhanced_prompt}\n")
 
 def generate_variations(prompt, num_variants=8, use_enhanced=True):
     instructions = [
-        "Create a photorealistic description for a detailed txt2img prompt in English: ",
-        "Write a creative, realistic, and detailed text-to-image prompt in English: ",
-        "Generate a descriptive and true to life txt2img prompt in English: ",
-        "Describe a photorealistic scene with detailed illumination for a txt2img prompt: ",
-        "Give a realistic, elegant txt2img prompt in English, emphasizing photorealism: ",
-        "Create a visually dynamic and hyperrealistic txt2img prompt in English: ",
-        "Write a cinematic txt2img prompt in English with hyperrealistic elements: ",
-        "Make a lifelike txt2img prompt in English, focusing on photorealistic depth: "
+        "Photorealistic description for txt2img prompt: ",
+        "Creative, realistic text-to-image prompt: ",
+        "Descriptive, true-to-life txt2img prompt: ",
+        "Naturalistic scene with detailed prompt: ",
+        "Realistic, elegant txt2img prompt: ",
+        "Visually dynamic, hyperrealistic prompt: ",
+        "Cinematic txt2img with hyperrealistic elements: ",
+        "Lifelike txt2img, focusing on photorealistic depth: "
     ]
-    
     if use_enhanced:
-        prompts = []
-        for i in range(num_variants):
-            instruction = instructions[i % len(instructions)]
-            enhanced_prompt = enhance_prompt(f"{instruction}{prompt}")
-            prompts.append(enhanced_prompt)
-    else:
+        prompts = [enhance_prompt(f"{instructions[i % len(instructions)]}{prompt}") for i in range(num_variants)]
+    else: 
         prompts = [prompt] * num_variants
     return prompts
 
+def generate_image(prompt, width, height, seed, model_name, client=client):
+    if not client:
+        st.error("No Hugging Face client available")
+        return None, seed, None
+
+    try:
+        with st.spinner("Generando imagen..."):
+            seed = int(seed) if seed != -1 else random.randint(0, AppConfig.MAX_SEED)
+            enhanced_prompt = enhance_prompt(prompt)
+            image = client.text_to_image(
+                prompt=enhanced_prompt, 
+                height=height, 
+                width=width, 
+                model=model_name, 
+                seed=seed
+            )
+            return image, seed, enhanced_prompt
+    except Exception as e:
+        st.error(f"Image generation error: {e}")
+        return None, seed, None
+
 def gen(prompts, width, height, model_name, num_variants=8):
     images = []
     seeds = []
-    
     while len(seeds) < num_variants:
         seed = random.randint(0, MAX_SEED)
-        if seed not in seeds:
+        if seed not in seeds: 
             seeds.append(seed)
     
     for i in range(num_variants):
         current_prompt = prompts[i] if len(prompts) > i else prompts[-1]
         with st.spinner(f"Generando imagen {i+1}/{num_variants}"):
-            image, used_seed = generate_image(current_prompt, width, height, seeds[i], model_name)
+            image, used_seed, enhanced_prompt = generate_image(current_prompt, width, height, seeds[i], model_name)
             if image:
                 image_path = DATA_PATH / f"generated_image_{used_seed}.jpg"
                 image.save(image_path)
-                save_prompt(f"generated_image_{used_seed}.jpg", current_prompt)
-                images.append(str(image_path))
+                save_prompt(f"generated_image_{used_seed}.jpg", enhanced_prompt)
+                images.append((str(image_path), enhanced_prompt))
                 st.success(f"Imagen {i+1} generada")
-    
     return images
 
 def get_prompt_for_image(image_name):
@@ -178,20 +120,41 @@ def get_prompt_for_image(image_name):
         with open(DATA_PATH / "prompts.txt", "r") as f:
             for line in f:
                 if line.startswith(image_name):
-                    return line.split(": ", 1)[1].strip()
-    except FileNotFoundError:
+                    original_prompt = line.split(": ", 1)[1].strip()
+                    image_path = DATA_PATH / image_name
+                    
+                    if image_path.exists():
+                        return original_prompt
+    except FileNotFoundError: 
         return "No hay prompt asociado"
     return "No hay prompt asociado"
+    
+def display_gallery():
+    st.header("Galería de Imágenes Guardadas")
+    images = list_saved_images()
+    
+    if images:
+        cols = st.columns(4) 
+        for i, image_file in enumerate(images):
+            with cols[i % 4]:
+                st.image(str(image_file), use_column_width=True)
+                prompt = get_prompt_for_image(image_file.name)
+                st.caption(prompt[:250])
+                if st.button(f"Borrar", key=f"delete_{i}_{image_file}"):
+                    if image_file.exists():
+                        os.remove(image_file)
+                        st.success("Imagen borrada")
+                        st.rerun()
 
 def login_form():
     st.title("Iniciar Sesión")
     username = st.text_input("Usuario", value="admin")
-    password = st.text_input("Contraseña", value="", type="password")
+    password = st.text_input("Contraseña", value="flux3x", type="password")
     if st.button("Iniciar Sesión"):
         if authenticate_user(username, password):
             st.session_state['authenticated'] = True
             st.success("Autenticación exitosa.")
-        else:
+        else: 
             st.error("Credenciales incorrectas. Intenta de nuevo.")
 
 def upload_image_to_gallery():
@@ -203,19 +166,20 @@ def upload_image_to_gallery():
         save_prompt(uploaded_image.name, "uploaded by user")
         st.sidebar.success(f"Imagen subida: {image_path}")
 
-app, swapper = prepare_face_app()
-
 def main():
-    if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
+    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')
+    
+    st.title("Flux +Upscale +Prompt Enhancer")
+    
+    if not client:
+        st.error("No se pudo establecer conexión con Hugging Face. Verifique sus tokens.")
+        return
+    
     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)
@@ -223,16 +187,17 @@ def main():
     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)
+            generated_images = gen(prompts, width, height, model_option, num_variants)
             
+            st.header("Imágenes Generadas")
+            cols = st.columns(4)
+            for i, (image_path, image_prompt) in enumerate(generated_images):
+                with cols[i % 4]:
+                    st.image(image_path, use_column_width=True)
+                    st.caption(image_prompt)
+    
     upload_image_to_gallery()
     display_gallery()
 
-if __name__ == "__main__":
+if __name__ == "__main__": 
     main()