Update app.py

app.py

@@ -1,24 +1,23 @@
 import os
-import gradio as gr
 import numpy as np
 import random
 from pathlib import Path
 from PIL import Image
-from huggingface_hub import AsyncInferenceClient, InferenceClient
+import streamlit as st
+from huggingface_hub import InferenceClient, AsyncInferenceClient
 from gradio_client import Client, handle_file
-from gradio_imageslider import ImageSlider
+import asyncio
 
 MAX_SEED = np.iinfo(np.int32).max
 HF_TOKEN = os.environ.get("HF_TOKEN")
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
 client = AsyncInferenceClient()
 llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-
 DATA_PATH = Path("./data")
-DATA_PATH.mkdir(exist_ok=True)  
+DATA_PATH.mkdir(exist_ok=True)
 
 def enable_lora(lora_add, basemodel):
-    return basemodel if not lora_add else lora_add
+    return lora_add if lora_add else basemodel
 
 async def generate_image(combined_prompt, model, width, height, scales, steps, seed):
     try:
@@ -56,91 +55,106 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
         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)  
+    
     if isinstance(image, str) and image.startswith("Error"):
-        return [image, None]
+        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)
 
     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")
-            return [image_path, DATA_PATH / f"upscale_image_{seed}.jpg"]
+            progress_bar.progress(100)
+            image_path.unlink()
+            return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
         else:
-            return [image_path, image_path]
+            progress_bar.empty()
+            return [str(image_path), str(prompt_file_path)]
     else:
-        return [image_path, image_path]
-
+        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 img2vid prompt in a single paragraph of up to 200 characters maximum, developing atmosphere, characters, lighting, and cameras.")
+        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.")
         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()
-
         return 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()
-    ]
-    usage = sum([f['size'] for f in files])
-    return [file["name"] for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
-
-css = """#col-container{ margin: 0 auto; max-width: 1024px;}"""
-
-with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as demo:
-    with gr.Column(elem_id="col-container"):
-        with gr.Row():
-            with gr.Column(scale=3):
-                output_res = ImageSlider(label="Generadas / Escaladas")
-            with gr.Column(scale=2):
-                prompt = gr.Textbox(label="Descripción de imagen")
-                basemodel_choice = gr.Dropdown(
-                    label="Modelo", 
-                    choices=["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-DEV"],
-                    value="black-forest-labs/FLUX.1-schnell"
-                )
-                lora_model_choice = gr.Dropdown(
-                    label="LORA Realismo", 
-                    choices=["Shakker-Labs/FLUX.1-dev-LoRA-add-details", "XLabs-AI/flux-RealismLora"],
-                    value="XLabs-AI/flux-RealismLora"
-                )
-                with gr.Row():
-                    process_lora = gr.Checkbox(label="Procesar LORA")
-                    process_upscale = gr.Checkbox(label="Procesar Escalador")
-                improved_prompt = gr.Textbox(label="Prompt Mejorado", interactive=False)
-                improve_btn = gr.Button("Mejorar prompt")
-                improve_btn.click(fn=improve_prompt, inputs=[prompt], outputs=improved_prompt)
-                with gr.Accordion(label="Opciones Avanzadas", open=False):
-                    width = gr.Slider(label="Ancho", minimum=512, maximum=1280, step=8, value=1280)
-                    height = gr.Slider(label="Alto", minimum=512, maximum=1280, step=8, value=768)
-                    upscale_factor = gr.Radio(label="Factor de Escala", choices=[2, 4, 8], value=2)
-                    scales = gr.Slider(label="Escalado", minimum=1, maximum=20, step=1, value=10)
-                    steps = gr.Slider(label="Pasos", minimum=1, maximum=100, step=1, value=20)
-                    seed = gr.Number(label="Semilla", value=-1)
-                btn = gr.Button("Generar")
-                btn.click(
-                    fn=gen, 
-                    inputs=[prompt, basemodel_choice, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model_choice, process_lora], 
-                    outputs=output_res
-                )
-        with gr.Row():
-            with gr.Column():
-                file_list = gr.Gallery(label="Imágenes Guardadas")  
-                storage_info = gr.Text(label="Uso de Almacenamiento")
-                refresh_btn = gr.Button("Actualizar Galería")
-                refresh_btn.click(fn=get_storage, inputs=None, outputs=[file_list, storage_info])
-                
-    demo.launch(allowed_paths=[str(DATA_PATH)])
+    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"
+
+def get_prompts():
+    prompt_files = list(DATA_PATH.glob("*.txt"))
+    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):
+    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))
+
+st.set_page_config(layout="wide")
+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"])
+format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
+process_lora = st.sidebar.checkbox("Procesar LORA")
+process_upscale = st.sidebar.checkbox("Procesar Escalador")
+
+if format_option == "9:16":
+    width = st.sidebar.slider("Ancho", 512, 720, 720, step=8)
+    height = st.sidebar.slider("Alto", 912, 1280, 1280, step=8)
+else:
+    width = st.sidebar.slider("Ancho", 512, 1280, 1280, step=8)
+    height = st.sidebar.slider("Alto", 512, 720, 720, step=8)
+
+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 st.sidebar.button("Mejorar prompt"):
+    improved_prompt = asyncio.run(improve_prompt(prompt))
+    st.session_state.improved_prompt = improved_prompt
+    st.write(f"{improved_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)
+    
+        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}")
+
+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(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}")