# Instalar las dependencias necesarias:
# Asegúrate de incluir estas en el archivo `requirements.txt` para Hugging Face Spaces:
# torch
# torchvision
# torchaudio
# diffusers
# huggingface_hub
# gradio

import torch
from diffusers import DiffusionPipeline
from huggingface_hub import login
import gradio as gr

# Configuración del modelo y adaptador LoRA
model_id = "camenduru/FLUX.1-dev"
adapter_id = "Emuixom/Trasher_Riddick"

# Cargar el pipeline y los pesos LoRA
pipeline = DiffusionPipeline.from_pretrained(model_id)
pipeline.load_lora_weights(adapter_id)

# Configuración del dispositivo (usar GPU si está disponible)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline.to(device)

# Función para generar imágenes
def generate_image(prompt, steps=15, width=512, height=512, guidance_scale=3.5, seed=1641421826):
    # Generar la imagen con el prompt proporcionado
    generator = torch.Generator(device=device).manual_seed(seed)
    image = pipeline(
        prompt=prompt,
        num_inference_steps=steps,
        generator=generator,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
    ).images[0]

    # Guardar la imagen generada temporalmente
    output_path = "output.png"
    image.save(output_path, format="PNG")
    return output_path

# Configurar la interfaz de Gradio
def gradio_interface(prompt, steps, width, height, guidance_scale, seed):
    output_path = generate_image(prompt, steps, width, height, guidance_scale, seed)
    return output_path

# Crear la interfaz
interface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe la imagen que deseas generar"),
        gr.Slider(label="Steps", minimum=1, maximum=50, value=15, step=1),
        gr.Slider(label="Width", minimum=128, maximum=1024, value=512, step=64),
        gr.Slider(label="Height", minimum=128, maximum=1024, value=512, step=64),
        gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=3.5, step=0.5),
        gr.Number(label="Seed", value=1641421826),
    ],
    outputs=gr.Image(type="file", label="Imagen Generada"),
    title="Generador de Imágenes con Diffusion y LoRA",
    description="Introduce un texto descriptivo para generar una imagen utilizando un modelo Diffusion con pesos LoRA.",
)

# Lanzar la aplicación
if __name__ == "__main__":
    interface.launch()