from pathlib import Path from PIL import Image import streamlit as st 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) def authenticate_user(username, password): return username == credentials["username"] and password == credentials["password"] async def gen(prompts, width, height, model_name, num_variants=1, use_enhanced=True): images = [] try: for idx, prompt in enumerate(prompts): seed = random.randint(0, MAX_SEED) image, seed = await generate_image(prompt, width, height, seed, model_name) image_path = save_image(image, f"generated_image_{seed}.jpg", prompt) if image_path: st.success(f"Imagen {idx + 1} generada") images.append(str(image_path)) except Exception as e: st.error(f"Error al generar imágenes: {e}") return images def list_saved_images(): return sorted(DATA_PATH.glob("*.jpg"), key=os.path.getmtime, reverse=True) def display_gallery(): 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"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") def run_async(func, *args): return asyncio.run(func(*args)) async def improve_prompt(prompt): try: instructions = [ "With this words, create a photorealistic description for a detailed txt2img prompt in English, 300 characters max.", "With this idea, write a creative, realistic, and detailed text-to-image prompt in English, 300 characters max.", "With this text, generate a descriptive and True to life txt2img prompt in English, 300 characters max.", "With my idea, describe a photorealistic scene with detailed illumination for a txt2img prompt in English, 300 characters max.", "With this concept, give a realistic, elegant txt2img prompt in English, emphasizing photorealism, 300 characters max.", "With this perspective, conform a visually dynamic and hyperrealistic txt2img prompt in English, 300 characters max.", "With this inspiration, realize a cinematic txt2img prompt in English with hyperrealistic elements, 300 characters max.", "With my idea, make a lifelike and txt2img prompt in English, focusing on photorealistic depth, 300 characters max." ] instruction = random.choice(instructions) formatted_prompt = f"{prompt}: {instruction}" response = llm_client.text_generation(formatted_prompt, max_new_tokens=100) return response['generated_text'][:100] 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 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 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.") async def generate_variations(prompt, num_variants, use_enhanced, style): prompts = set() while len(prompts) < num_variants: if use_enhanced: enhanced_prompt = await improve_prompt(f"{prompt}, estilo: {style}") prompts.add(enhanced_prompt) else: prompts.add(f"{prompt}, estilo: {style}") return list(prompts) async def main(): st.set_page_config(layout="wide") if 'authenticated' not in st.session_state or not st.session_state['authenticated']: login_form() return prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=900) format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9", "1:1"]) model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"]) prompt_checkbox = st.sidebar.checkbox("Mejorar Prompt") style_type = st.sidebar.selectbox("Estilo", ["Realismo", "Hiperrealismo", "Photorealismo", "Ilustración", "Abstracto", "Imaginativo"]) style_option = st.sidebar.selectbox("Camara", ["Full Shot", "Over Showder", "Medium Shot", "Long Shot", "Close-UP", "Dutch Angle", "High Angle", "Low Angle", "Oblique Angle"]) width, height = (360, 640) if format_option == "9:16" else (640, 360) if format_option == "16:9" else (640, 640) if prompt_checkbox: num_variants = st.sidebar.slider("Número de imágenes a generar", 1, 8, 1) else: num_variants = 1 if prompt_checkbox: with st.spinner("Generando prompts mejorados..."): prompts = await generate_variations(f"{prompt}, estilo: {style_option}, {style_type}", num_variants, prompt_checkbox, style_option) if prompt_checkbox else [f"{prompt}, estilo: {style_option}, {style_type}"] else: prompts = [f"{prompt}, estilo: {style_option}"] if st.sidebar.button("Generar Imágenes"): with st.spinner("Generando imágenes..."): try: results = await gen(prompts, width, height, model_option, num_variants, prompt_checkbox) st.session_state['generated_image_paths'] = results for result in results: st.image(result, caption="Imagen Generada") except Exception as e: st.error(f"Error al generar las imágenes: {str(e)}") display_gallery() if __name__ == "__main__": asyncio.run(main())