Update app.py

app.py

@@ -1,8 +1,8 @@
 from pathlib import Path
 from PIL import Image
 import streamlit as st
-import os, random, numpy as np, yaml, time, 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
 
@@ -21,16 +21,28 @@ class AppConfig:
     CLEANUP_DAYS: int = 7    
 
 MAX_SEED = AppConfig.MAX_SEED
-client = InferenceClient()
+
+HF_TOKENS = [
+    os.environ.get("HF_TOKEN1"),
+    os.environ.get("HF_TOKEN2"),
+    os.environ.get("HF_TOKEN3")
+]
+
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
-HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN")
 
-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 get_inference_client(tokens=HF_TOKENS):
+    for token in tokens:
+        try:
+            client = InferenceClient(token=token)
+            client.text_generation("Test", max_length=10)
+            return client
+        except Exception as e:
+            st.warning(f"Token fallido: {token[:5]}...")
+    st.error("Todos los tokens de HuggingFace han fallado")
+    return None
+
+client = get_inference_client()
 
 def authenticate_user(username, password): 
     return username == credentials["username"] and password == credentials["password"]
@@ -38,15 +50,36 @@ def authenticate_user(username, password):
 def list_saved_images(): 
     return sorted(DATA_PATH.glob("*.jpg"), key=lambda x: x.stat().st_mtime, reverse=True)
 
-def enhance_prompt(text):
+def get_image_description(image_path, original_prompt):
+    if not client:
+        return original_prompt
+    
+    try:
+        description_prompt = f"Analyze this image in great detail. The original prompt was: {original_prompt}. Provide a comprehensive and creative description focusing on visual elements, mood, and artistic style."
+        
+        description = client.text_generation(
+            description_prompt,
+            model="mistralai/Mixtral-8x7B-v0.1",
+            max_new_tokens=150
+        )
+        
+        return description.strip()
+    except Exception as e:
+        st.warning(f"Image description generation error: {e}")
+        return original_prompt
+
+def enhance_prompt(text, client=client):
+    if not client:
+        return text[:200]
     try:
         enhanced = client.text_generation(
-            "With this text, generate a descriptive and photorealistic txt2img prompt in English in 200 characters maximum: " + text,
+            "Generate a photorealistic, detailed txt2img prompt in 200 characters maximum: " + text,
             model="mistralai/Mixtral-8x7B-v0.1", 
             max_length=200
         )
         return enhanced[:200]
-    except: 
+    except Exception as e:
+        st.warning(f"Prompt enhancement error: {e}")
         return text[:200]
 
 def save_prompt(image_name, prompt):
@@ -55,14 +88,14 @@ def save_prompt(image_name, prompt):
 
 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: ",
+        "Creative, realistic text-to-image prompt: ",
+        "Descriptive, true-to-life txt2img prompt: ",
+        "Photorealistic scene with detailed illumination: ",
+        "Realistic, elegant txt2img prompt: ",
+        "Visually dynamic, hyperrealistic prompt: ",
+        "Cinematic txt2img with hyperrealistic elements: ",
+        "Lifelike txt2img, focusing on photorealistic depth: "
     ]
     if use_enhanced:
         prompts = [enhance_prompt(f"{instructions[i % len(instructions)]}{prompt}") for i in range(num_variants)]
@@ -70,15 +103,25 @@ def generate_variations(prompt, num_variants=8, use_enhanced=True):
         prompts = [prompt] * num_variants
     return prompts
 
-def generate_image(prompt, width, height, seed, model_name):
+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)
+            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"Error al generar imagen: {e}")
+        st.error(f"Image generation error: {e}")
         return None, seed, None
 
 def gen(prompts, width, height, model_name, num_variants=8):
@@ -101,6 +144,22 @@ def gen(prompts, width, height, model_name, num_variants=8):
                 st.success(f"Imagen {i+1} generada")
     return images
 
+def get_prompt_for_image(image_name):
+    try:
+        with open(DATA_PATH / "prompts.txt", "r") as f:
+            for line in f:
+                if line.startswith(image_name):
+                    original_prompt = line.split(": ", 1)[1].strip()
+                    image_path = DATA_PATH / image_name
+                    
+                    if image_path.exists():
+                        return get_image_description(str(image_path), original_prompt)
+                    
+                    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()
@@ -111,27 +170,17 @@ def display_gallery():
             with cols[i % 4]:
                 st.image(str(image_file), use_column_width=True)
                 prompt = get_prompt_for_image(image_file.name)
-                st.caption(prompt[:100])
+                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 get_prompt_for_image(image_name):
-    try:
-        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: 
-        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="flux3x", type="password")
+    password = st.text_input("Contraseña", value="", type="password")
     if st.button("Iniciar Sesión"):
         if authenticate_user(username, password):
             st.session_state['authenticated'] = True
@@ -155,6 +204,10 @@ def main():
     
     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"])
@@ -168,7 +221,6 @@ def main():
         if st.sidebar.button("Generar Imágenes"):
             generated_images = gen(prompts, width, height, model_option, num_variants)
             
-            # Display generated images with their prompts
             st.header("Imágenes Generadas")
             cols = st.columns(4)
             for i, (image_path, image_prompt) in enumerate(generated_images):