Update app.py

app.py

@@ -57,10 +57,10 @@ 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, process_enhancer):
+async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
     combined_prompt = prompt  
     if process_enhancer:
-        improved_prompt = await improve_prompt(prompt)
+        improved_prompt = await improve_prompt(prompt, language)
         combined_prompt = f"{prompt} {improved_prompt}"
     
     if seed == -1:
@@ -92,11 +92,12 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
         progress_bar.progress(100)
         return [str(image_path), str(prompt_file_path)]
 
-async def improve_prompt(prompt):
+async def improve_prompt(prompt, language):
     try:
         instruction_en = "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters..."
         instruction_es = "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_en} {instruction_es}"
+        instruction = instruction_en if language == "en" else instruction_es
+        formatted_prompt = f"{prompt}: {instruction}"
         response = llm_client.text_generation(formatted_prompt, max_new_tokens=300)
         improved_text = response['generated_text'].strip() if 'generated_text' in response else response.strip()
         return improved_text[:300] if len(improved_text) > 300 else improved_text
@@ -145,6 +146,7 @@ def main():
     scales = st.sidebar.slider("Escalado", 1, 20, 10)
     steps = st.sidebar.slider("Pasos", 1, 100, 20)
     seed = st.sidebar.number_input("Semilla", value=-1)
+    language = st.sidebar.selectbox("Idioma", ["es", "en"])
 
     if format_option == "9:16":
         width = 720
@@ -155,7 +157,7 @@ def main():
 
     if st.sidebar.button("Generar Imagen"):
         with st.spinner("Mejorando y generando imagen..."):
-            result = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer))
+            result = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language))
             image_paths = result[0]
             prompt_file = result[1]
 
@@ -197,4 +199,4 @@ def main():
                     st.error(f"Error al borrar la imagen o prompt: {e}")
 
 if __name__ == "__main__":
-    main()
+    main()