app6.py

import gradio as gr
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# --- Configuración del Modelo (Sencilla) ---

# Forzamos el dispositivo a CPU (para evitar errores de CUDA/VRAM)
device = "cpu"
dtype_config = torch.float32 

# Modelo Image-to-Image LucyEdit
model_id = "kpsss34/LucyEdit_hybrid" 

pipe = None
try:
    # Carga sencilla, sin device_map, sin aceleración avanzada
    pipe = DiffusionPipeline.from_pretrained(
        model_id, 
        torch_dtype=dtype_config, 
        use_safetensors=True
    ).to(device) 
    
    print("✅ Modelo LucyEdit cargado con éxito en la CPU.")

except Exception as e:
    print(f"❌ Error CRÍTICO al cargar el modelo: {e}")
    print("El modelo NO ha podido cargarse. Comprueba tu conexión o memoria RAM.")


# --- Función de Procesamiento con Difusión (i2i) ---
def procesar_con_difusion(imagen_entrada, estilo_radial):
    """
    Modifica la imagen usando el pipeline de difusión con el estilo radial seleccionado.
    """
    # Manejo de error de carga
    if pipe is None:
        # Devuelve el error si el modelo no se cargó correctamente
        return Image.new('RGB', (512, 512), color = 'red') 

    if imagen_entrada is None:
        return None 
    
    # 1. Prompt Base Fijo (Lo que debe hacer el modelo)
    prompt_base = "fotografía de alta calidad, retrato detallado, con estilo" 
    
    estilo_prompts = {
        # El prompt fijo + el modificador del estilo
        "Blanco y Negro (Monocromático)": ", monocromático, alto contraste, dramático, blanco y negro",
        "Alto Contraste y Saturación": ", colores vívidos, alto contraste, HDR, saturación extrema",
        "Original (Poco Ruido)": ", fotografía de alta calidad, realista, colores naturales, sutil",
    }
    
    full_prompt = prompt_base + estilo_prompts.get(estilo_radial, "")

    # 2. Fuerza de Difusión FIJA
    # Un valor alto (0.85) para asegurar que el efecto (B/N o contraste) se aplique bien.
    STRENGTH_FIJA = 0.85 
    
    # 3. Preprocesar la imagen (LucyEdit trabaja en 512x512)
    init_image = imagen_entrada.convert("RGB").resize((512, 512))
    
    try:
        # 4. Ejecutar el pipeline de difusión i2i
        image = pipe(
            prompt=full_prompt, 
            image=init_image, 
            strength=STRENGTH_FIJA, 
            guidance_scale=7.5    
        ).images[0] 
        
        return image
    except Exception as e:
        print(f"Error durante la ejecución del pipeline: {e}")
        # Devuelve un cuadro de error si el proceso falla
        return Image.new('RGB', (512, 512), color = 'red')


# --- Interfaz Gradio con gr.Blocks() ---
with gr.Blocks(title="LucyEdit Estilos Fijos") as demo:
    gr.Markdown(
        """
        # 🖼️ Tarea de Difusión (Image-to-Image) con LucyEdit
        Carga una imagen y selecciona un **Estilo Radial** (Blanco y Negro o Contraste) para que el modelo de difusión la transforme. 
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil", 
                label="1. Cargar Imagen Inicial",
            )
            
            estilo_radial = gr.Radio(
                ["Original (Poco Ruido)", "Blanco y Negro (Monocromático)", "Alto Contraste y Saturación"],
                label="2. Selecciona el Estilo de Transformación",
                value="Blanco y Negro (Monocromático)"
            )
            
            process_button = gr.Button("✨ Aplicar Difusión", variant="primary")
            
        with gr.Column(scale=1):
            image_output = gr.Image(
                type="pil",
                label="Imagen Transformada por Difusión",
                height=512 
            )
    
    process_button.click(
        fn=procesar_con_difusion, 
        inputs=[image_input, estilo_radial],
        outputs=image_output
    )

demo.launch(inbrowser=True)