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

# --- Configuración del Modelo (Forzado a CPU) ---

# Establecemos el dispositivo a CPU ya que CUDA no está disponible en tu entorno.
device = "cpu"
dtype_config = torch.float32 # Los modelos en CPU usan float32

# Usamos Stable Diffusion v1.5, el modelo más ligero que funciona bien.
model_id = "runwayml/stable-diffusion-v1-5" 

pipe = None
try:
    # Carga sin accelerate (sin device_map) y forzando a CPU.
    # Esto elimina las dependencias problemáticas de tu error.
    pipe = DiffusionPipeline.from_pretrained(
        model_id, 
        torch_dtype=dtype_config, 
        use_safetensors=True
    ).to(device) # <--- Forzamos la carga a la CPU
    
    print("✅ Modelo Stable Diffusion v1.5 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
    prompt_base = "fotografía de alta calidad, retrato detallado, foto con textura" 
    
    estilo_prompts = {
        "Blanco y Negro (Monocromático)": ", monocromático, alto contraste, película de 35mm, grain, dramático",
        "Alto Contraste y Saturación": ", colores vívidos, alto contraste, saturación extrema, cinematográfico, hyperdetailed",
        "Original (Poco Ruido)": ", fotografía de alta calidad, realista, colores naturales, sutil, cinematic lighting",
    }
    
    full_prompt = prompt_base + estilo_prompts.get(estilo_radial, "")

    # 2. Fuerza de Difusión FIJA (strength): 
    STRENGTH_FIJA = 0.9 
    
    # 3. Preprocesar la imagen (SD v1.5 nativo es 512x512)
    # Convertimos a RGB y redimensionamos
    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="SD v1.5 Estilos Fijos") as demo:
    gr.Markdown(
        """
        # ✅ Tarea de Difusión (Image-to-Image) con SD v1.5
        Carga una imagen y selecciona un **Estilo Radial** para que el modelo de difusión la transforme. 
        **NOTA:** La generación será lenta (varios minutos) ya que se ejecuta en CPU.
        """
    )
    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)