Update app.py

app.py

@@ -36,16 +36,24 @@ def get_upscale_finegrain(prompt, img_path, upscale_factor):
     try:
         client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
         result = client.predict(
-            input_image=handle_file(img_path), prompt=prompt, negative_prompt="",
-            seed=42, upscale_factor=upscale_factor, controlnet_scale=0.6,
-            controlnet_decay=1, condition_scale=6, tile_width=112,
-            tile_height=144, denoise_strength=0.35, num_inference_steps=18,
-            solver="DDIM", api_name="/process"
+            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:
         return None
 
+# Función para guardar el prompt
+def save_prompt(prompt_text, seed):
+    try:
+        prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
+        with open(prompt_file_path, "w") as prompt_file:
+            prompt_file.write(prompt_text)
+        return prompt_file_path
+    except Exception as e:
+        st.error(f"Error al guardar el prompt: {e}")
+        return None
+
+# Función principal de generación de imágenes
 async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora):
     model = enable_lora(lora_model, basemodel) if process_lora else basemodel
     improved_prompt = await improve_prompt(prompt)
@@ -62,12 +70,8 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
         progress_bar.empty()
         return [image, None, combined_prompt]
 
-    image_path = DATA_PATH / f"image_{seed}.jpg"
-    image.save(image_path, format="JPEG")
-
-    prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
-    with open(prompt_file_path, "w") as prompt_file:
-        prompt_file.write(combined_prompt)
+    image_path = save_image(image, seed)
+    prompt_file_path = save_prompt(combined_prompt, seed)
 
     if process_upscale:
         upscale_image_path = get_upscale_finegrain(combined_prompt, image_path, upscale_factor)
@@ -75,7 +79,7 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
             upscale_image = Image.open(upscale_image_path)
             upscale_image.save(DATA_PATH / f"upscale_image_{seed}.jpg", format="JPEG")
             progress_bar.progress(100)
-            image_path.unlink()
+            image_path.unlink()  # Borrar la imagen original si se escaló
             return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
         else:
             progress_bar.empty()
@@ -86,22 +90,16 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
 
 async def improve_prompt(prompt):
     try:
-        instruction = ("With this idea, describe in English a detailed txt2img prompt in a single paragraph of up to 100 characters maximum, developing atmosphere, characters, lighting, and cameras.")
+        instruction = ("With this idea, describe in English a detailed txt2img prompt...")
         formatted_prompt = f"{prompt}: {instruction}"
         response = llm_client.text_generation(formatted_prompt, max_new_tokens=200)
         improved_text = response['generated_text'].strip() if 'generated_text' in response else response.strip()
-        
-        if len(improved_text) > 200:
-            improved_text = improved_text[:200]
-        
-        return improved_text
+        return improved_text[:200] if len(improved_text) > 200 else improved_text
     except Exception as e:
         return f"Error mejorando el prompt: {e}"
 
 def get_storage():
-    files = [{"name": str(file.resolve()), "size": file.stat().st_size}
-             for file in DATA_PATH.glob("*.jpg") if file.is_file()]
-    
+    files = [{"name": str(file.resolve()), "size": file.stat().st_size} for file in DATA_PATH.glob("*.jpg") if file.is_file()]
     usage = sum([f['size'] for f in files])
     return [f["name"] for f in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
 
@@ -109,12 +107,6 @@ def get_prompts():
     prompt_files = [file for file in DATA_PATH.glob("*.txt") if file.is_file()]
     return {file.stem.replace("prompt_", ""): file for file in prompt_files}
 
-def run_gen():
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    prompt_to_use = st.session_state.get('improved_prompt', prompt)
-    return loop.run_until_complete(gen(prompt_to_use, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora))
-
 def delete_image(image_path):
     try:
         if Path(image_path).exists():
@@ -125,13 +117,21 @@ def delete_image(image_path):
     except Exception as e:
         st.error(f"Error al borrar la imagen: {e}")
 
+def run_gen():
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    prompt_to_use = st.session_state.get('improved_prompt', prompt)
+    return loop.run_until_complete(gen(prompt_to_use, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora))
+
+# Configuración de la página y sidebar
 st.set_page_config(layout="wide")
 st.title("Generador de Imágenes FLUX")
 prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=200)
 
+# Opciones avanzadas
 with st.sidebar.expander("Opciones avanzadas", expanded=False):
-    basemodel = st.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-DEV"])
-    lora_model = st.selectbox("LORA Realismo", ["Shakker-Labs/FLUX.1-dev-LoRA-add-details", "XLabs-AI/flux-RealismLora"])
+    basemodel = st.selectbox("Modelo Base", ["FLUX.1", "FLUX.1-DEV"])
+    lora_model = st.selectbox("LORA Realismo", ["FLUX.1-dev-LoRA", "RealismLora"])
     format_option = st.selectbox("Formato", ["9:16", "16:9"])
     process_lora = st.checkbox("Procesar LORA")
     process_upscale = st.checkbox("Procesar Escalador")
@@ -150,7 +150,7 @@ else:
 if st.sidebar.button("Mejorar prompt"):
     improved_prompt = asyncio.run(improve_prompt(prompt))
     st.session_state.improved_prompt = improved_prompt
-    st.write(f"{improved_prompt}")
+    st.write(f"Prompt mejorado: {improved_prompt}")
 
 if st.sidebar.button("Generar Imagen"):
     with st.spinner("Generando imagen..."):
@@ -182,9 +182,15 @@ for idx, file in enumerate(files):
         image = Image.open(file)
         prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
         prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
+        
         st.image(image, caption=f"Imagen {idx+1}")
         st.write(f"Prompt: {prompt_text}")
+        
         if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
-            os.remove(file)
-            if prompt_file:
-                os.remove(prompt_file)
+            try:
+                os.remove(file)  
+                if prompt_file:
+                    os.remove(prompt_file)  
+                st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
+            except Exception as e:
+                st.error(f"Error al borrar la imagen o prompt: {e}")