Update app.py

app.py

@@ -10,7 +10,8 @@ import asyncio
 from concurrent.futures import ThreadPoolExecutor
 
 MAX_SEED = np.iinfo(np.int32).max
-HF_TOKEN = os.environ.get("HF_TOKEN_UPSCALER")
+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")
@@ -23,9 +24,6 @@ def run_async(func):
     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
-
 async def generate_image(combined_prompt, model, width, height, scales, steps, seed):
     try:
         if seed == -1:
@@ -59,10 +57,8 @@ 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
+async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer):
     combined_prompt = prompt  
-
     if process_enhancer:
         improved_prompt = await improve_prompt(prompt)
         combined_prompt = f"{prompt} {improved_prompt}"
@@ -71,7 +67,7 @@ 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)
-    image, seed = await generate_image(combined_prompt, model, width, height, scales, steps, seed)
+    image, seed = await generate_image(combined_prompt, basemodel, width, height, scales, steps, seed)
     progress_bar.progress(50)
 
     if isinstance(image, str) and image.startswith("Error"):
@@ -98,11 +94,12 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
 
 async def improve_prompt(prompt):
     try:
-        instruction = ("With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add ilumination, admosphere, cinematic and characters...")
-        formatted_prompt = f"{prompt}: {instruction}"
-        response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
+        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}"
+        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[:500] if len(improved_text) > 500 else improved_text
+        return improved_text[:300] if len(improved_text) > 300 else improved_text
     except Exception as e:
         return f"Error mejorando el prompt: {e}"
 
@@ -137,15 +134,13 @@ def delete_image(image_path):
 
 def main():
     st.set_page_config(layout="wide")
-    st.title("FLUX +prompt/enhancer +upscaler +LORA")
-    
-    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
-    process_enhancer = st.sidebar.checkbox("Mejorar Prompt", value=False) 
+    st.title("FLUX with prompt enhancer and upscaler")
+
+    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=200)
+    process_enhancer = st.sidebar.checkbox("Mejorar Prompt", value=True)
     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=False)
-    process_upscale = st.sidebar.checkbox("Procesar Escalador", value=False)
+    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)
@@ -160,7 +155,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, lora_model, process_lora, process_enhancer))
+            result = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer))
             image_paths = result[0]
             prompt_file = result[1]