Update app.py

app.py

@@ -7,6 +7,7 @@ import streamlit as st
 from huggingface_hub import InferenceClient, AsyncInferenceClient
 from gradio_client import Client, handle_file
 import asyncio
+from concurrent.futures import ThreadPoolExecutor
 
 MAX_SEED = np.iinfo(np.int32).max
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -16,6 +17,13 @@ llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 
+def run_async(func):
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    executor = ThreadPoolExecutor(max_workers=1)
+    result = loop.run_in_executor(executor, func)
+    return loop.run_until_complete(result)
+
 def enable_lora(lora_add, basemodel):
     return lora_add if lora_add else basemodel
 
@@ -52,10 +60,13 @@ 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):
+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
-    improved_prompt = await improve_prompt(prompt)
-    combined_prompt = f"{prompt} {improved_prompt}"
+    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)
@@ -86,9 +97,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)]
 
-async def improve_prompt(prompt):
+async def improve_prompt(prompt, language):
     try:
-        instruction = ("With this idea, describe in English a detailed txt2img prompt in 300 characters at most...")
+        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...")
+
         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()
@@ -96,51 +111,61 @@ async def improve_prompt(prompt):
     except Exception as e:
         return f"Error mejorando el prompt: {e}"
 
-def save_image(image, seed):
-    try:
-        image_path = DATA_PATH / f"image_{seed}.jpg"
-        image.save(image_path, format="JPEG")
-        return image_path
-    except Exception as e:
-        st.error(f"Error al guardar la imagen: {e}")
-        return None
+async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora, process_enhancer, prompt_language):
+    model = enable_lora(lora_model, basemodel) if process_lora else basemodel
+    combined_prompt = prompt  # Usar el prompt original por defecto
 
-def get_storage():
-    files = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
-    files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
-    usage = sum([file.stat().st_size for file in files])
-    return [str(file.resolve()) for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
+    # 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}"
+    
+    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]
 
-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}
+    image_path = save_image(image, seed)
+    prompt_file_path = save_prompt(combined_prompt, seed)
 
-def delete_image(image_path):
-    try:
-        if Path(image_path).exists():
-            Path(image_path).unlink()
-            st.success(f"Imagen {image_path} borrada.")
+    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:
-            st.error("El archivo de imagen no existe.")
-    except Exception as e:
-        st.error(f"Error al borrar la imagen: {e}")
+            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 main():
     st.set_page_config(layout="wide")
-    st.title("FLUX with enhancer upscaler and LORA realism training")
-    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
-
-    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"])
-        format_option = st.selectbox("Formato", ["9:16", "16:9"])
-        process_lora = st.checkbox("Procesar LORA")
-        process_upscale = st.checkbox("Procesar Escalador")
-        upscale_factor = st.selectbox("Factor de Escala", [2, 4, 8], index=0)
-        scales = st.slider("Escalado", 1, 20, 10)
-        steps = st.slider("Pasos", 1, 100, 20)
-        seed = st.number_input("Semilla", value=-1)
+    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
+    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"])
+    process_lora = st.sidebar.checkbox("Procesar LORA", value=True)
+    process_upscale = st.sidebar.checkbox("Procesar Escalador", value=True)
+    upscale_factor = st.sidebar.selectbox("Factor de Escala", [2, 4, 8], index=0)
+    scales = st.sidebar.slider("Escalado", 1, 20, 10)
+    steps = st.sidebar.slider("Pasos", 1, 100, 20)
+    seed = st.sidebar.number_input("Semilla", value=-1)
 
     if format_option == "9:16":
         width = 720
@@ -151,11 +176,7 @@ def main():
 
     if st.sidebar.button("Generar Imagen"):
         with st.spinner("Mejorando y generando imagen..."):
-            improved_prompt = asyncio.run(improve_prompt(prompt))
-            st.session_state.improved_prompt = improved_prompt
-            prompt_to_use = st.session_state.get('improved_prompt', prompt)
-            result = asyncio.run(gen(prompt_to_use, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora))
-
+            result = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora, process_enhancer, prompt_language))
             image_paths = result[0]
             prompt_file = result[1]
 
@@ -172,29 +193,5 @@ def main():
                 st.write(f"Prompt utilizado: {prompt_text}")
             else:
                 st.write("El archivo del prompt no está disponible.")
-
-    files, usage = get_storage()
-    st.text(usage)
-    cols = st.columns(6)
-    prompts = get_prompts()
-
-    for idx, file in enumerate(files):
-        with cols[idx % 6]:
-            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}"):
-                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}")
-
-if __name__ == "__main__":
-    main()
+        else:
+            st.error("No se pudo generar la imagen.")