Update app.py

app.py

@@ -25,12 +25,16 @@ DATA_PATH.mkdir(exist_ok=True)
 PREDEFINED_SEED = random.randint(0, MAX_SEED)
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
 
+if not HF_TOKEN_UPSCALER:
+    st.warning("HF_TOKEN_UPSCALER no está configurado. Algunas funcionalidades pueden no funcionar.")
+
 def get_upscale_finegrain(prompt, img_path, upscale_factor):
     try:
-        upscale_client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
+        upscale_client = InferenceClient("finegrain/finegrain-image-enhancer", 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:
+    except Exception as e:
+        st.error(f"Error al mejorar la imagen: {e}")
         return None
 
 def authenticate_user(username, password):
@@ -70,11 +74,19 @@ async def generate_image(prompt, width, height, seed, model_name):
     image = await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
     return image, seed
 
-async def gen(prompt, width, height, model_name):
-    seed = PREDEFINED_SEED
-    image, seed = await generate_image(prompt, width, height, seed, model_name)
-    image_path = save_image(image, f"generated_image_{seed}.jpg", prompt)
-    return str(image_path)
+async def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=True):
+    images = []
+    try:
+        for idx, prompt in enumerate(prompts):
+            seed = random.randint(0, MAX_SEED)
+            image, seed = await generate_image(prompt, width, height, seed, model_name)
+            image_path = save_image(image, f"generated_image_{seed}.jpg")
+            if image_path:
+                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 list_saved_images():
     return list(DATA_PATH.glob("*.jpg"))
@@ -95,9 +107,12 @@ def display_gallery():
                     st.success("Imagen seleccionada")
 
                 if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"):
-                    os.remove(image_file)
-                    st.success("Imagen borrada")
-                    display_gallery()
+                    if os.path.exists(image_file):
+                        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.")
 
@@ -111,23 +126,32 @@ def run_async(func, *args):
 
 async def improve_prompt(prompt):
     try:
-        instruction = ("With this idea, describe in English a detailed txt2img prompt in 500 characters maximun")
+        instructions = [
+            "With my idea create a vibrant description for a detailed txt2img prompt, 300 characters max.",
+            "With my idea write a creative and detailed text-to-image prompt in English, 300 characters max.",
+            "With my idea generate a descriptive and visual txt2img prompt in English, 300 characters max.",
+            "With my idea describe a photorealistic with illumination txt2img prompt in English, 300 characters max.",
+            "With my idea give a realistic and elegant txt2img prompt in English, 300 characters max.",
+            "With my idea conform a visually dynamic and surreal txt2img prompt in English, 300 characters max.",
+            "With my idea realize an artistic and cinematic txt2img prompt in English, 300 characters max.",
+            "With my idea make a narrative and immersive txt2img prompt in English, 300 characters max."
+        ]
+        instruction = random.choice(instructions)
         formatted_prompt = f"{prompt}: {instruction}"
-        response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
-        return response['generated_text'][:500] if 'generated_text' in response else response.strip()
+        response = llm_client.text_generation(formatted_prompt, max_new_tokens=100)
+        return response['generated_text'][:100] if 'generated_text' in response else response.strip()
     except Exception as e:
         return f"Error mejorando el prompt: {e}"
 
-def save_image(image, file_name, prompt=None):
-    image_path = DATA_PATH / file_name
-    if image_path.exists():
-        st.warning(f"La imagen '{file_name}' ya existe en la galería. No se guardó.")
-        return None 
-    else:
-        image.save(image_path, format="JPEG")
-        if prompt:
-            save_prompt(f"{file_name}: {prompt}")  
-        return image_path
+async def generate_variations(prompt, num_variants, use_enhanced):
+    prompts = set() 
+    while len(prompts) < num_variants:
+        if use_enhanced:
+            enhanced_prompt = await improve_prompt(prompt)
+            prompts.add(enhanced_prompt)
+        else:
+            prompts.add(prompt)
+    return list(prompts)
 
 def get_prompt_for_image(image_name):
     prompts = {}
@@ -152,6 +176,15 @@ def login_form():
         else:
             st.error("Credenciales incorrectas. Intenta de nuevo.")
 
+def save_image(image, filename):
+    try:
+        image_path = DATA_PATH / filename
+        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:
@@ -176,27 +209,27 @@ async def main():
     model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"])
     prompt_checkbox = st.sidebar.checkbox("Prompt Enhancer")
     upscale_checkbox = st.sidebar.checkbox("Escalar imagen")
-    width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
-    upload_image_to_gallery()
+    width, height = (360, 640) if format_option == "9:16" else (640, 360) if format_option == "16:9" else (640, 640)
+    
+    num_variants = st.sidebar.slider("Número de imágenes a generar", 1, 8, 1) if prompt_checkbox else 1
     
     if prompt_checkbox:
-        with st.spinner("Mejorando el prompt..."):
-            try:
-                improved_prompt = run_async(improve_prompt, prompt)
-            except Exception as e:
-                st.error(f"Error al mejorar el prompt: {str(e)}")
-                improved_prompt = prompt
+        with st.spinner("Generando prompts mejorados..."):
+            prompts = await generate_variations(prompt, num_variants, True)
     else:
-        improved_prompt = prompt
+        prompts = [prompt] 
+
+    upload_image_to_gallery()
 
-    if st.sidebar.button("Generar Imagen"):
-        with st.spinner("Generando imagen..."):
+    if st.sidebar.button("Generar Imágenes"):
+        with st.spinner("Generando imágenes..."):
             try:
-                result = run_async(gen, improved_prompt, width, height, model_option)  # Usar el improved_prompt
-                st.session_state['generated_image_path'] = result
-                st.image(result, caption="Imagen Generada")
+                results = await gen(prompts, width, height, model_option, num_variants, prompt_checkbox)
+                st.session_state['generated_image_paths'] = results
+                for result in results:
+                    st.image(result, caption="Imagen Generada")
             except Exception as e:
-                st.error(f"Error al generar la imagen: {str(e)}")
+                st.error(f"Error al generar las imágenes: {str(e)}")
 
     if generated_image_path:
         if upscale_checkbox: