from pathlib import Path from PIL import Image import streamlit as st import insightface from insightface.app import FaceAnalysis from huggingface_hub import InferenceClient, AsyncInferenceClient import asyncio import os import random import numpy as np import yaml try: with open("config.yaml", "r") as file: credentials = yaml.safe_load(file) except Exception as e: st.error(f"Error al cargar el archivo de configuración: {e}") credentials = {"username": "", "password": ""} MAX_SEED = np.iinfo(np.int32).max client = AsyncInferenceClient() llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1") DATA_PATH = Path("./data") DATA_PATH.mkdir(exist_ok=True) PREDEFINED_SEED = random.randint(0, MAX_SEED) HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER") def get_upscale_finegrain(prompt, img_path, upscale_factor): try: upscale_client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER) result = upscale_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 except Exception: return None def authenticate_user(username, password): return username == credentials["username"] and password == credentials["password"] def prepare_face_app(): app = FaceAnalysis(name='buffalo_l') app.prepare(ctx_id=0, det_size=(640, 640)) swapper = insightface.model_zoo.get_model('onix.onnx') return app, swapper app, swapper = prepare_face_app() def sort_faces(faces): return sorted(faces, key=lambda x: x.bbox[0]) def get_face(faces, face_id): if not faces or len(faces) < face_id: raise ValueError("Rostro no disponible.") return faces[face_id - 1] def swap_faces(source_image, source_face_index, destination_image, destination_face_index): faces = sort_faces(app.get(source_image)) source_face = get_face(faces, source_face_index) res_faces = sort_faces(app.get(destination_image)) if destination_face_index > len(res_faces) or destination_face_index < 1: raise ValueError("Índice de rostro de destino no válido.") res_face = get_face(res_faces, destination_face_index) result = swapper.get(destination_image, res_face, source_face, paste_back=True) return result async def generate_image(prompt, width, height, seed, model_name): if seed == -1: seed = random.randint(0, MAX_SEED) image = await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name) return image, seed async def gen(prompt, width, height, model_name): seed = PREDEFINED_SEED image, seed = await generate_image(prompt, width, height, seed, model_name) image_path = save_image(image, f"generated_image_{seed}.jpg", prompt) return str(image_path) def list_saved_images(): return list(DATA_PATH.glob("*.jpg")) def display_gallery(): st.header("Galería de Imágenes Guardadas") images = list_saved_images() if images: cols = st.columns(8) for i, image_file in enumerate(images): with cols[i % 8]: st.image(str(image_file), caption=image_file.name, use_column_width=True) prompt = get_prompt_for_image(image_file.name) st.write(prompt[:300]) if st.button(f"Usar", key=f"select_{i}_{image_file.name}"): st.session_state['generated_image_path'] = str(image_file) st.success("Imagen seleccionada") if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"): os.remove(image_file) st.success("Imagen borrada") display_gallery() else: st.info("No hay imágenes guardadas.") def save_prompt(prompt): with open(DATA_PATH / "prompts.txt", "a") as f: f.write(prompt + "\n") st.success("Prompt guardado.") def run_async(func, *args): return asyncio.run(func(*args)) async def improve_prompt(prompt): try: instruction = ("With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters if needed...") formatted_prompt = f"{prompt}: {instruction}" response = llm_client.text_generation(formatted_prompt, max_new_tokens=500) return response['generated_text'][:500] if 'generated_text' in response else response.strip() except Exception as e: return f"Error mejorando el prompt: {e}" def save_image(image, file_name, prompt=None): image_path = DATA_PATH / file_name if image_path.exists(): st.warning(f"La imagen '{file_name}' ya existe en la galería. No se guardó.") return None else: image.save(image_path, format="JPEG") if prompt: save_prompt(f"{file_name}: {prompt}") return image_path def get_prompt_for_image(image_name): prompts = {} try: with open(DATA_PATH / "prompts.txt", "r") as f: for line in f: if line.startswith(image_name): prompts[image_name] = line.split(": ", 1)[1].strip() except FileNotFoundError: return "No hay prompt asociado." return prompts.get(image_name, "No hay prompt asociado.") def login_form(): st.title("Iniciar Sesión") username = st.text_input("Usuario", value="admin") password = st.text_input("Contraseña", value="flux3x", type="password") if st.button("Iniciar Sesión"): if authenticate_user(username, password): st.success("Autenticación exitosa.") st.session_state['authenticated'] = True else: st.error("Credenciales incorrectas. Intenta de nuevo.") def upload_image_to_gallery(): uploaded_image = st.file_uploader("Sube una imagen a la galería", type=["jpg", "jpeg", "png"]) if uploaded_image: image = Image.open(uploaded_image) image_path = save_image(image, f"{uploaded_image.name}") if image_path: save_prompt("uploaded by user") st.success(f"Imagen subida: {image_path}") def main(): st.set_page_config(layout="wide") if 'authenticated' not in st.session_state or not st.session_state['authenticated']: login_form() return st.title("Flux +Upscale +Prompt Enhancer +FaceSwap") generated_image_path = st.session_state.get('generated_image_path') st.header("Generador de Imágenes") prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500) format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"]) model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"]) prompt_checkbox = st.sidebar.checkbox("Prompt Enhancer") upscale_checkbox = st.sidebar.checkbox("Escalar imagen") width, height = (720, 1280) if format_option == "9:16" else (1280, 720) upload_image_to_gallery() if prompt_checkbox: with st.spinner("Mejorando el prompt..."): try: improved_prompt = run_async(improve_prompt, prompt) except Exception as e: st.error(f"Error al mejorar el prompt: {str(e)}") improved_prompt = prompt else: improved_prompt = prompt if st.sidebar.button("Generar Imagen"): with st.spinner("Generando imagen..."): try: result = run_async(gen, improved_prompt, width, height, model_option) # Usar el improved_prompt st.session_state['generated_image_path'] = result st.image(result, caption="Imagen Generada") except Exception as e: st.error(f"Error al generar la imagen: {str(e)}") if generated_image_path: if upscale_checkbox: with st.spinner("Escalando imagen..."): try: upscale_image_path = get_upscale_finegrain("Upscale", generated_image_path, 2) if upscale_image_path: st.image(upscale_image_path, caption="Imagen Escalada") except Exception as e: st.error(f"Error al escalar la imagen: {str(e)}") st.header("Intercambio de Rostros") source_image_file = st.file_uploader("Imagen de Origen", type=["jpg", "jpeg", "png"]) if source_image_file is not None: try: source_image = Image.open(source_image_file) except Exception as e: st.error(f"Error al cargar la imagen de origen: {str(e)}") source_image = None else: source_image = Image.open("face.jpg") source_face_index = st.number_input('Posición del Rostro', min_value=1, value=1, key="source_face_index") destination_face_index = st.number_input('Posición del Rostro de Destino', min_value=1, value=1, key="destination_face_index") if st.button("Intercambiar Rostros"): try: destination_image = Image.open(generated_image_path) result_image = swap_faces(np.array(source_image), source_face_index, np.array(destination_image), destination_face_index) swapped_image = Image.fromarray(result_image) swapped_image_path = save_image(swapped_image, f"swapped_image_{PREDEFINED_SEED}.jpg") if swapped_image_path: st.image(swapped_image, caption="Intercambio de Rostro") os.remove(generated_image_path) else: st.warning("La imagen intercambiada ya existe en la galería.") except Exception as e: st.error(f"Ocurrió un error al intercambiar rostros: {str(e)}") display_gallery() if __name__ == "__main__": main()