Update app.py

app.py

@@ -1,131 +1,125 @@
-import os
-import random
-import asyncio
+import os, numpy as np, random, asyncio, insightface
 from pathlib import Path
 from PIL import Image
 from insightface.app import FaceAnalysis
 import streamlit as st
-from huggingface_hub import InferenceClient, AsyncInferenceClient
+from huggingface_hub import InferenceClient
+from gradio_client import Client, handle_file
+from concurrent.futures import ThreadPoolExecutor
 import yaml
-import insightface
 
-try:
-    with open("config.yaml", "r") as file:
-        credentials = yaml.safe_load(file)
-except Exception:
-    credentials = {"username": "", "password": ""}
-
-MAX_SEED = 2**31 - 1
+MAX_SEED = np.iinfo(np.int32).max
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
-client = AsyncInferenceClient()
-llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-DATA_PATH = Path("./data").mkdir(exist_ok=True)
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+DATA_PATH = Path("./data")
+DATA_PATH.mkdir(exist_ok=True)
 
 def prepare_face_app():
     app = FaceAnalysis(name='buffalo_l')
     app.prepare(ctx_id=0, det_size=(640, 640))
-    return app, insightface.model_zoo.get_model('onix.onnx')
+    swapper = insightface.model_zoo.get_model('onix.onnx')
+    return app, swapper
 
 app, swapper = prepare_face_app()
 
-def run_async(func):
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    return loop.run_until_complete(loop.run_in_executor(None, func))
-
-async def generate_image(prompt, model, width, height, scales, steps, seed):
+async def generate_image(combined_prompt, model, width, height, scales, steps, seed):
     if seed == -1:
         seed = random.randint(0, MAX_SEED)
-    try:
-        return await client.text_to_image(prompt=prompt, height=height, width=width, guidance_scale=scales, num_inference_steps=steps, model=model), seed
-    except Exception as e:
-        return f"Error: {e}", None
-
-def upscale_image(image_path, factor):
-    try:
-        image = Image.open(image_path)
-        new_size = (image.width * factor, image.height * factor)
-        return image.resize(new_size, Image.LANCZOS)
-    except Exception as e:
-        st.error(f"Error: {e}")
-
-def save_image(image, seed):
-    try:
-        path = DATA_PATH / f"image_{seed}.jpg"
-        image.save(path, format="JPEG")
-        return path
-    except Exception as e:
-        st.error(f"Error: {e}")
-
-def save_prompt(prompt, seed):
-    try:
-        path = DATA_PATH / f"prompt_{seed}.txt"
-        with open(path, 'w') as f:
-            f.write(prompt)
-        return path
-    except Exception as e:
-        st.error(f"Error: {e}")
-
-async def gen(prompt, model, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
-    combined_prompt = f"{prompt} {await improve_prompt(prompt, language)}" if process_enhancer else prompt
-    if seed == -1: seed = random.randint(0, MAX_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"): return [image, None]
+    seed = int(seed)
+    image = await client.text_to_image(prompt=combined_prompt, height=height, width=width, guidance_scale=scales, num_inference_steps=steps, model=model)
+    return image, seed
+
+def get_upscale_finegrain(prompt, img_path, upscale_factor):
+    client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
+    result = client.predict(input_image=handle_file(img_path), prompt=prompt, upscale_factor=upscale_factor)
+    return result[1] if isinstance(result, list) and len(result) > 1 else None
+
+def save_prompt(prompt_text, seed):
+    prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
+    with open(prompt_file_path, "w") as prompt_file:
+        prompt_file.write(prompt_text)
+    return prompt_file_path
+
+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, language)
+        combined_prompt = f"{prompt} {improved_prompt}"
+    
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    image, seed = await generate_image(combined_prompt, basemodel, width, height, scales, steps, seed)
     image_path = save_image(image, seed)
-    prompt_path = save_prompt(combined_prompt, seed)  
-    if process_upscale: image = upscale_image(image_path, upscale_factor)
-    progress_bar.progress(100)
-    return [str(image_path) if not process_upscale else str(save_image(image, seed)), str(prompt_path)]
+    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")
+            image_path.unlink()  
+            return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
+        else:
+            return [str(image_path), str(prompt_file_path)]
+    else:
+        return [str(image_path), str(prompt_file_path)]
 
 async def improve_prompt(prompt, language):
-    instruction = "With this idea, describe in English a detailed txt2img prompt..." if language == "en" else "Con esta idea, describe en español un prompt detallado..."
+    instruction = "With this idea, describe in English a detailed txt2img prompt in 500 characters..." if language == "en" else "Con esta idea, describe en español un prompt detallado de txt2img en 500 caracteres..."
     formatted_prompt = f"{prompt}: {instruction}"
-    response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
-    return response.get('generated_text', '').strip()[:500]
-
-def authenticate_user(username, password, credentials):
-    return username == credentials["username"] and password == credentials["password"]
-
-def login_form(credentials):
-    username, password = st.text_input("Usuario"), st.text_input("Contraseña", type="password")
-    if st.button("Iniciar Sesión"):
-        if authenticate_user(username, password, credentials):
-            st.session_state['authenticated'] = True
-            st.session_state['username'] = username  
-        else:
-            st.error("Credenciales incorrectas")
+    response = client.text_generation(formatted_prompt, max_new_tokens=500)
+    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
+
+def save_image(image, seed):
+    image_path = DATA_PATH / f"image_{seed}.jpg"
+    image.save(image_path, format="JPEG")
+    return image_path
+
+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)
+    return [str(file.resolve()) for file in files]
+
+def delete_image(image_path):
+    if Path(image_path).exists():
+        Path(image_path).unlink()
 
 def main():
     st.set_page_config(layout="wide")
+    
     if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
-        st.title("Iniciar Sesión")
-        login_form(credentials)
+        st.warning("Por favor, inicia sesión para acceder a la aplicación.")
         return
-
-    st.sidebar.text_input("Descripción de la imagen", max_chars=900, key='prompt')
+    
+    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=900)
     process_enhancer = st.sidebar.checkbox("Mejorar Prompt", value=False)
     language = st.sidebar.selectbox("Idioma", ["en", "es"])
     basemodel = st.sidebar.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-DEV", "black-forest-labs/FLUX.1-schnell"])
     format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
     process_upscale = st.sidebar.checkbox("Procesar Escalador", value=False)
     upscale_factor = st.sidebar.selectbox("Factor de Escala", [2, 4, 8], index=0)
-    scales, steps, seed = 7.5, 50, -1
-    
+    scales = st.sidebar.slider("Escalado", 1, 20, 10)
+    steps = st.sidebar.slider("Pasos", 1, 100, 20)
+    seed = st.sidebar.number_input("Semilla", value=-1)
+    width, height = (1080, 1920) if format_option == "9:16" else (1920, 1080)
+
     if st.sidebar.button("Generar Imagen"):
-        result = run_async(gen(st.session_state['prompt'], basemodel, 512, 512, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language))
-        if result[0] and isinstance(result[0], str):
-            st.image(result[0], caption='Imagen Generada')
-            st.write(f"Prompt: {result[1]}") 
-            if st.button("Borrar Imagen"):
-                os.remove(result[0]) 
-                st.success("Imagen borrada.")
-            if st.button("Face Swap"):
-                st.write("Funcionalidad de Face Swap no implementada")  
-    
-    if st.sidebar.button("Cerrar Sesión"):
-        st.session_state.clear()  
+        with st.spinner("Generando..."):
+            image_path, prompt_file_path = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language))
+            if image_path:
+                st.image(image_path, caption="Imagen Generada", use_column_width=True)
+                st.download_button("Descargar Imagen", image_path)
+
+    files = get_storage()
+    for file in files:
+        st.image(file)
+        st.text(f"Prompt: {file.stem}")
+        if st.button(f"Swap Face {file.stem}"):
+            pass
+        if st.button(f"Borrar {file.stem}"):
+            delete_image(file)
 
 if __name__ == "__main__":
     main()