Update app.py

app.py

@@ -9,8 +9,12 @@ import random
 import numpy as np
 import yaml
 
-with open("config.yaml", "r") as file:
-    credentials = yaml.safe_load(file)
+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 = InferenceClient()
@@ -19,17 +23,28 @@ DATA_PATH.mkdir(exist_ok=True)
 PREDEFINED_SEED = random.randint(0, MAX_SEED)
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN")
 
+if not HF_TOKEN_UPSCALER:
+    st.warning("HF_TOKEN no está configurado. Algunas funcionalidades pueden no funcionar.")
+
 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
+    try:
+        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
+    except Exception as e:
+        st.error(f"Error al mejorar la imagen: {e}")
+        return 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
+    try:
+        image_files = sorted(DATA_PATH.glob("*.jpg"))
+        return image_files
+    except Exception as e:
+        st.error(f"Error al listar imágenes guardadas: {e}")
+        return []
 
 def prepare_face_app():
     app = FaceAnalysis(name='buffalo_l')
@@ -43,29 +58,31 @@ 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))
-    if not faces:
-        st.warning("No se encontraron caras en la imagen de origen.")
-        return None
     source_face = get_face(faces, source_face_index)
-    
     res_faces = sort_faces(app.get(destination_image))
-    if not res_faces:
-        st.warning("No se encontraron caras en la imagen de destino.")
-        return None
+    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):
     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
+    try:
+        image = client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
+        return image, seed
+    except Exception as e:
+        st.error(f"Error al generar imagen: {e}")
+        if hasattr(e, 'response') and e.response is not None:
+            st.error(f"Detalles del error: {e.response.text}") 
+        return None, seed
 
 def display_gallery():
     st.header("Galería de Imágenes Guardadas")
@@ -84,67 +101,74 @@ def display_gallery():
                     if image_file.exists():
                         os.remove(image_file)
                         st.success("Imagen borrada")
+                        display_gallery()
                     else:
                         st.warning("La imagen no existe.")
     else:
         st.info("No hay imágenes guardadas.")
 
-def get_prompt_for_image(image_name):
-    prompts = {}
-    prompt_file = DATA_PATH / "prompts.txt"
-    if prompt_file.exists():
-        with open(prompt_file, "r") as f:
-            for line in f:
-                if line.startswith(image_name):
-                    prompts[image_name] = line.split(": ", 1)[1].strip()
-    return prompts.get(image_name, "No hay prompt asociado.")
+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):
-    if not use_enhanced:
+    if 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)
+    else:
         return [prompt] * num_variants
-    
-    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)):
-        clean_prompt = prompt.split(": ", 1)[-1] 
-        image, _ = generate_image(clean_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: {clean_prompt}")
-            st.success(f"Imagen {idx + 1} generada")
-            images.append(str(image_path))
+    seeds = []
+    
+    while len(seeds) < num_variants:
+        new_seed = random.randint(0, MAX_SEED)
+        if new_seed not in seeds: 
+            seeds.append(new_seed)
+            
+    try:
+        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))
+    except Exception as e:
+        st.error(f"Error al generar imágenes: {e}")
     return images
 
-def save_prompt(prompt):
-    clean_prompt = prompt.split(": ", 1)[-1]  
-    with open(DATA_PATH / "prompts.txt", "a") as f:
-        f.write(clean_prompt + "\n")  
-        st.success("Prompt guardado.")
+def get_prompt_for_image(image_name):
+    prompts = {}
+    try:
+        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") 
+    password = st.text_input("Contraseña", value="flux3", type="password") 
     if st.button("Iniciar Sesión"):
         if authenticate_user(username, password):
             st.success("Autenticación exitosa.")
@@ -153,24 +177,42 @@ def login_form():
             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
+    try:
+        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
+    except Exception as e:
+        st.error(f"Error al guardar la imagen: {e}")
+        return None
 
 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)
-        unique_filename = f"{random.randint(100000, 999999)}_{uploaded_image.name}"
-        image_path = save_image(image, unique_filename)
+        image_path = save_image(image, f"{uploaded_image.name}")
         if image_path:
-            save_prompt(f"{unique_filename}: uploaded by user")
+            save_prompt(f"{uploaded_image.name}: uploaded by user")
             st.sidebar.success(f"Imagen subida: {image_path}")
 
+def gen(prompts, width, height, model_name, num_variants=1):
+    images = []
+    try:
+        for idx, prompt in enumerate(prompts[:num_variants]):
+            seed = random.randint(0, MAX_SEED)
+            image, seed = 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))
+    except Exception as e:
+        st.error(f"Error al generar imágenes: {e}")
+    return images
+
 def main():
     st.set_page_config(layout="wide")
     if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
@@ -185,16 +227,20 @@ def main():
     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 or st.sidebar.button("Generar Imágenes", key="force_generate"):
+    
+    if prompt:
         prompts = generate_variations(prompt, num_variants=num_variants, use_enhanced=prompt_enhance)
-        images = gen(prompts, width, height, model_option, num_variants)
+        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()