app6.py

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

# --- Configuración del Modelo (Cambiado a Stable Diffusion v1.5) ---

# SD v1.5 requiere menos VRAM (4-8 GB) y es más estable en entornos limitados.
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype_config = torch.float16 if device == "cuda" else torch.float32 

# Modelo Stable Diffusion v1.5
model_id = "runwayml/stable-diffusion-v1-5" 

pipe = None
try:
    # Carga con accelerate (device_map="auto") para manejar la memoria
    pipe = DiffusionPipeline.from_pretrained(
        model_id, 
        torch_dtype=dtype_config, 
        use_safetensors=True,
        device_map="auto" # Mantiene accelerate para optimizar
    )
    print("Modelo Stable Diffusion v1.5 cargado con éxito usando accelerate.")

except Exception as e:
    print(f"Error CRÍTICO al cargar el modelo: {e}")
    print("El modelo NO ha podido cargarse. Si esto ocurre con v1.5, el hardware es insuficiente.")


# --- Función de Procesamiento con Difusión (i2i) ---
# Hemos quitado 'strength_slider' de los argumentos
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:
        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 = {
        # El prompt fijo + el modificador del estilo
        "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): 
    # Es fundamental para que el proceso i2i aplique el estilo del prompt.
    STRENGTH_FIJA = 0.9 
    
    # 3. Preprocesar la imagen (SD v1.5 nativo es 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, # Usamos la fuerza fija
            guidance_scale=7.5    
        ).images[0] 
        
        return image
    except Exception as e:
        print(f"Error durante la ejecución del pipeline: {e}")
        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. 
        Este modelo es más ligero y evita la imagen roja.
        """
    )
    with gr.Row():
        # Lado izquierdo: Inputs y Controles
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil", 
                label="1. Cargar Imagen Inicial",
            )
            
            # --- Control Radial (Radio Buttons) ---
            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)"
            )
            # -----------------------------------
            
            # ELIMINAMOS EL SLIDER DE FUERZA DE DIFUSIÓN
            
            process_button = gr.Button("✨ Aplicar Difusión", variant="primary")
            
        # Lado derecho: Output
        with gr.Column(scale=1):
            image_output = gr.Image(
                type="pil",
                label="Imagen Transformada por Difusión",
                height=512 
            )
    
    # Conexión de la acción: ¡Cambiamos los inputs a solo dos!
    process_button.click(
        fn=procesar_con_difusion, 
        inputs=[image_input, estilo_radial],
        outputs=image_output
    )

demo.launch(inbrowser=True)