Update app.py

app.py

@@ -54,19 +54,18 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
     if seed == -1:
         seed = random.randint(0, MAX_SEED)
     seed = int(seed)
-    
     progress_bar = st.progress(0)
-    progress_bar.progress(10) 
     image, seed = await generate_image(combined_prompt, model, width, height, scales, steps, seed)
-    progress_bar.progress(50)  
-    
+    progress_bar.progress(50)
+
     if isinstance(image, str) and image.startswith("Error"):
         progress_bar.empty()
         return [image, None, combined_prompt]
 
     image_path = DATA_PATH / f"image_{seed}.jpg"
     image.save(image_path, format="JPEG")
-    
+
+    # Save the prompt used in a text file
     prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
     with open(prompt_file_path, "w") as prompt_file:
         prompt_file.write(combined_prompt)
@@ -83,9 +82,9 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
             progress_bar.empty()
             return [str(image_path), str(prompt_file_path)]
     else:
-        progress_bar.progress(100) 
+        progress_bar.progress(100)
         return [str(image_path), str(prompt_file_path)]
-        
+
 async def improve_prompt(prompt):
     try:
         instruction = ("With this idea, describe in English a detailed txt2img prompt in a single paragraph of up to 200 characters maximum, developing atmosphere, characters, lighting, and cameras.")
@@ -97,21 +96,25 @@ async def improve_prompt(prompt):
         return f"Error mejorando el prompt: {e}"
 
 def get_storage():
-    files = list(DATA_PATH.glob("*.jpg"))
-    usage = sum(file.stat().st_size for file in files)
-    return files, f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
+    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 [file["name"] for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
 
 def get_prompts():
-    prompt_files = list(DATA_PATH.glob("*.txt"))
+    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(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora):
+def run_gen():
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
-    return loop.run_until_complete(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora))
+    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))
 
 st.set_page_config(layout="wide")
-st.title("Generador de Imágenes FLUX")
+st.title("Generador de Imágenes FLUX y Escalador con IA")
+
 prompt = st.sidebar.text_input("Descripción de la imagen")
 basemodel = st.sidebar.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-DEV"])
 lora_model = st.sidebar.selectbox("LORA Realismo", ["Shakker-Labs/FLUX.1-dev-LoRA-add-details", "XLabs-AI/flux-RealismLora"])
@@ -138,12 +141,12 @@ if st.sidebar.button("Mejorar prompt"):
 
 if st.sidebar.button("Generar Imagen"):
     with st.spinner("Generando imagen..."):
-        image_paths, prompt_file = run_gen(st.session_state.get('improved_prompt', prompt), basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora)
+        image_paths, prompt_file = run_gen()
     
-        if image_paths and isinstance(image_paths[0], str) and Path(image_paths[0]).exists():
-            st.image(image_paths[0], caption="Imagen Generada")
-            prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
-            st.write(f"Prompt utilizado: {prompt_text}")
+    if image_paths:
+        st.image(image_paths[0], caption="Imagen Generada")
+        prompt_text = Path(prompt_file).read_text()
+        st.write(f"Prompt utilizado: {prompt_text}")
 
 files, usage = get_storage()
 st.text(usage)
@@ -153,10 +156,7 @@ prompts = get_prompts()
 for idx, file in enumerate(files):
     with cols[idx % 6]:
         image = Image.open(file)
-        prompt_file = prompts.get(file.stem.replace("image_", ""), None)
+        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}")
-        
-        if st.button(f"Eliminar Imagen {idx+1}", key=f"delete_{idx}"):
-            file.unlink()
-            st.experimental_rerun()
+        st.write(f"Prompt: {prompt_text}")