Update app.py

app.py

@@ -60,49 +60,18 @@ def save_prompt(prompt_text, seed):
         st.error(f"Error al guardar el prompt: {e}")
         return None
 
-async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora, process_enhancer):
-    model = enable_lora(lora_model, basemodel) if process_lora else basemodel
-    combined_prompt = prompt  
-
-    if process_enhancer:
-        improved_prompt = await improve_prompt(prompt)
-        combined_prompt = f"{prompt} {improved_prompt}"
-    
-    if seed == -1:
-        seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
-    progress_bar = st.progress(0)
-    image, seed = await generate_image(combined_prompt, model, width, height, scales, steps, seed)
-    progress_bar.progress(50)
-
-    if isinstance(image, str) and image.startswith("Error"):
-        progress_bar.empty()
-        return [image, None, 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)
-        if upscale_image_path:
-            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()  
-            return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
-        else:
-            progress_bar.empty()
-            return [str(image_path), str(prompt_file_path)]
-    else:
-        progress_bar.progress(100)
-        return [str(image_path), str(prompt_file_path)]
+def save_image(image, seed):
+    image_path = DATA_PATH / f"generated_image_{seed}.jpg"
+    image.save(image_path)
+    return image_path
 
 async def improve_prompt(prompt, language):
     try:
-        if language == "English":
-            instruction = ("With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters...")
-        else:  # Spanish
-            instruction = ("Con esta idea, describe en español un prompt detallado de txt2img en un máximo de 500 caracteres, añadiendo iluminación, atmósfera, elementos cinematográficos y personajes...")
+        instruction = (
+            "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters..."
+            if language == "English"
+            else "Con esta idea, describe en español un prompt detallado de txt2img en un máximo de 500 caracteres, añadiendo iluminación, atmósfera, elementos cinematográficos y personajes..."
+        )
 
         formatted_prompt = f"{prompt}: {instruction}"
         response = llm_client.text_generation(formatted_prompt, max_new_tokens=300)
@@ -115,7 +84,6 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
     model = enable_lora(lora_model, basemodel) if process_lora else basemodel
     combined_prompt = prompt  # Usar el prompt original por defecto
 
-    # Mejorar el prompt si el checkbox está activado y según el idioma seleccionado
     if process_enhancer:
         improved_prompt = await improve_prompt(prompt, prompt_language)
         combined_prompt = f"{prompt} {improved_prompt}"
@@ -149,14 +117,13 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
         progress_bar.progress(100)
         return [str(image_path), str(prompt_file_path)]
 
-
 def main():
     st.set_page_config(layout="wide")
     st.title("FLUX with enhancer and upscaler with LORA model training")
     
     prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=200)
     process_enhancer = st.sidebar.checkbox("Mejorar Prompt", value=True)  
-    prompt_language = st.sidebar.selectbox("Idioma para mejorar el prompt", ["English", "Spanish"])  # Selector de idioma
+    prompt_language = st.sidebar.selectbox("Idioma para mejorar el prompt", ["English", "Spanish"])
     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"])
     format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
@@ -194,4 +161,7 @@ def main():
             else:
                 st.write("El archivo del prompt no está disponible.")
         else:
-            st.error("No se pudo generar la imagen.")
+            st.error("No se pudo generar la imagen.")
+
+if __name__ == "__main__":
+    main()