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

# 1. Cargar el modelo ligero (Solo se ejecuta una vez al iniciar el Space)
model_id = "runwayml/stable-diffusion-v1-5" 
device = "cuda" if torch.cuda.is_available() else "cpu"

# Usamos StableDiffusionImg2ImgPipeline para Image-to-Image
try:
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe = pipe.to(device)
except:
    # Fallback si no hay GPU o si falla la carga con float16
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id)
    pipe = pipe.to(device)


# 2. Función de procesamiento (Aquí se ejecuta la inferencia)
def convert_to_bn_diffusion(input_image: Image.Image, prompt: str, strength: float) -> Image.Image:
    """
    Ejecuta el pipeline de difusión I2I para estilizar la imagen.
    """
    
    # Aseguramos que la imagen esté en RGB y redimensionamos si es necesario
    input_image = input_image.convert("RGB").resize((512, 512))
    
    # Prompt forzando el estilo monocromático
    bn_prompt = f"{prompt}, high contrast, black and white, monochrome, grayscale"

    # Generación (Inferencia)
    output_image = pipe(
        prompt=bn_prompt,
        image=input_image,
        strength=strength, # Cuanto más alto, más se transforma (más BN)
        guidance_scale=7.5
    ).images[0]
    
    return output_image

# 3. Interfaz de Gradio
iface = gr.Interface(
    fn=convert_to_bn_diffusion,
    inputs=[
        gr.Image(type="pil", label="Sube tu imagen (se redimensionará a 512x512)"),
        gr.Textbox(label="Prompt adicional (ej: 'a moody photo', 'vintage style')", value="A sharp, detailed photograph"),
        gr.Slider(minimum=0.5, maximum=1.0, step=0.05, value=0.9, label="Fuerza de Estilización (Strength - Cuánto se convierte a B/N)")
    ],
    outputs="image",
    title="Conversor a B/N con Modelo de Difusión (Img2Img)",
    description="Sube una imagen y el modelo Stable Diffusion intentará convertirla a blanco y negro basado en el prompt y la fuerza de estilización."
)

if __name__ == "__main__":
    iface.launch()