Update app.py

app.py

@@ -8,55 +8,76 @@ import os
 import random
 import numpy as np
 import yaml
+import time
+import logging
+from typing import Tuple, Optional, List, Dict
+from dataclasses import dataclass
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+@dataclass
+class AppConfig:
+    MAX_SEED: int = np.iinfo(np.int32).max
+    DEFAULT_WIDTH: int = 1280
+    DEFAULT_HEIGHT: int = 720
+    SUPPORTED_FORMATS: List[str] = ["9:16", "16:9", "1:1"]
+    CLEANUP_DAYS: int = 7
+    MAX_GALLERY_SIZE: int = 1000
 
 try:
     with open("config.yaml", "r") as file:
+        config = yaml.safe_load(file)
+except Exception as e:
+    logger.error(f"Error al cargar el archivo de configuración: {e}")
+    config = {
+        "development": {
+            "max_image_size": 1280,
+            "supported_formats": ["9:16", "16:9", "1:1"],
+            "models": ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"],
+            "storage": {"cleanup_after_days": 7, "max_gallery_size": 1000}
+        }
+    }
+
+try:
+    with open("credentials.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
+MAX_SEED = AppConfig.MAX_SEED
 client = InferenceClient()
 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")
 
-if not HF_TOKEN_UPSCALER:
-    st.warning("HF_TOKEN no está configurado. Algunas funcionalidades pueden no funcionar.")
+@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_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=handle_file(img_path), prompt=prompt, upscale_factor=upscale_factor)
+        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)
+
 def authenticate_user(username, password):
     return username == credentials["username"] and password == credentials["password"]
 
 def list_saved_images():
-    try:
-        image_files = sorted(
-            DATA_PATH.glob("*.jpg"),
-            key=lambda x: x.stat().st_mtime,
-            reverse=True
-        )
-        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')
-    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()
+    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])
@@ -78,19 +99,18 @@ def swap_faces(source_image, source_face_index, destination_image, destination_f
 
 def generate_image(prompt, width, height, seed, model_name):
     try:
-        seed = int(seed) if seed != -1 else random.randint(0, MAX_SEED)
-        image = client.text_to_image(
-            prompt=prompt,
-            height=height,
-            width=width,
-            model=model_name,
-            seed=seed
-        )
-        return image, seed
+        with st.spinner("Generando imagen..."):
+            seed = int(seed) if seed != -1 else random.randint(0, MAX_SEED)
+            image = client.text_to_image(
+                prompt=prompt,
+                height=height,
+                width=width,
+                model=model_name,
+                seed=seed
+            )
+            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():
@@ -111,119 +131,104 @@ def display_gallery():
                         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 save_prompt(prompt):
+def save_prompt(image_name, prompt):
     with open(DATA_PATH / "prompts.txt", "a") as f:
-        f.write(prompt + "\n")
-        st.success("Prompt guardado.")
+        f.write(f"{image_name}: {prompt[:100]}\n")
+
+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 generate_variations(prompt, num_variants, use_enhanced):
+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: "
+    ]
+    
     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)
+        prompts = []
+        for i in range(num_variants):
+            instruction = instructions[i % len(instructions)]
+            enhanced_prompt = enhance_prompt(f"{instruction}{prompt}")
+            prompts.append(enhanced_prompt)
     else:
-        return [prompt] * num_variants
+        prompts = [prompt] * num_variants
+    return prompts
 
-def gen(prompts, width, height, model_name, num_variants=1):
+def gen(prompts, width, height, model_name, num_variants=8):
     images = []
     seeds = []
     
     while len(seeds) < num_variants:
-        new_seed = random.randint(0, MAX_SEED)
-        if new_seed not in seeds:
-            seeds.append(new_seed)
+        seed = random.randint(0, MAX_SEED)
+        if seed not in seeds:
+            seeds.append(seed)
     
-    try:
-        prompts_to_use = prompts[:num_variants]
-        if len(prompts_to_use) < num_variants:
-            prompts_to_use.extend([prompts_to_use[-1]] * (num_variants - len(prompts_to_use)))
-        
-        for idx, (seed, prompt) in enumerate(zip(seeds, prompts_to_use)):
-            image, used_seed = generate_image(prompt, width, height, seed, model_name)
+    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)
             if image:
-                image_path = save_image(image, f"generated_image_{seed}.jpg")
-                if image_path:
-                    clean_prompt = prompt.split(": ")[-1] if ": " in prompt else prompt
-                    save_prompt(f"generated_image_{seed}.jpg: {clean_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}")
+                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))
+                st.success(f"Imagen {i+1} generada")
     
     return images
 
 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):
-                    full_prompt = line.split(": ", 1)[1].strip()
-                    if ":" in full_prompt:
-                        prompt = full_prompt.split(": ")[-1].strip()
-                    else:
-                        prompt = full_prompt
-                    prompts[image_name] = prompt
+                    return line.split(": ", 1)[1].strip()
     except FileNotFoundError:
-        return "No hay prompt asociado."
-    return prompts.get(image_name, "No hay prompt asociado.")
+        return "No hay prompt asociado"
+    return "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="flux3", 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.")
             st.session_state['authenticated'] = True
+            st.success("Autenticación exitosa.")
         else:
             st.error("Credenciales incorrectas. Intenta de nuevo.")
 
-def save_image(image, filename):
-    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)
-        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}")
+        image_path = DATA_PATH / uploaded_image.name
+        image.save(image_path)
+        save_prompt(uploaded_image.name, "uploaded by user")
+        st.sidebar.success(f"Imagen subida: {image_path}")
+
+app, swapper = prepare_face_app()
 
 def main():
     st.set_page_config(layout="wide")
     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')
     prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500)