import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Modelo base
modelo_base = "runwayml/stable-diffusion-v1-5"  # Puedes cambiarlo si usas otro modelo compatible

# Cargar el pipeline
pipe = StableDiffusionPipeline.from_pretrained(
    modelo_base,
    torch_dtype=torch.float16
).to("cpu")

# Cargar LoRA (ajusta a tu ruta)
pipe.load_lora_weights("DRDELATV/lora-image-caricature", weight_name="JNX_Funny_Caricature (1).safetensors")

# Función para generar la imagen
def generar(prompt):
    image = pipe(prompt).images[0]
    return image

# Interfaz con Gradio
demo = gr.Interface(
    fn=generar,
    inputs=gr.Textbox(label="Prompt de entrada", placeholder="Describe la imagen que quieres..."),
    outputs=gr.Image(type="pil", label="Imagen generada"),
    title="Generador de caricaturas con LoRA",
    description="Este generador usa Stable Diffusion + LoRA para crear imágenes caricaturescas."
)

demo.launch()