import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
import random
from PIL import Image
import os

# Configuración del modelo - Usando Pony XL NSFW realista
model_id = "John6666/sexoholic-real-pony-nsfw-v2-sdxl"
scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Términos para prompts en estilo Pony XL (basados en tags booru)
quality_boosters = ["score_9", "score_8_up", "score_7_up", "score_6_up", "score_5_up", "score_4_up"]
styles = ["photorealistic", "hyperrealistic", "ultra detailed", "highres", "masterpiece", "best quality"]
nsfw_tags = ["rating_explicit", "nude", "explicit nudity", "detailed genitalia", "sensual pose", "erotic"]
lighting = ["volumetric godrays", "cinematic lighting", "soft shadows", "dramatic lighting"]
camera = ["canon eos r5, 85mm lens", "close-up", "full body shot", "dynamic angle"]
settings = ["luxury bedroom interior", "urban cityscape at night", "tropical beach sunset", "modern art studio"]
negative_prompt_base = "score_4, score_5, score_6, blurry, deformed, ugly, poorly drawn, mutated, extra limbs, watermark, text, low quality"

# Función para generar prompts en estilo Pony XL
def generate_prompts(input_prompt, num_prompts=5):
    if not input_prompt:
        input_prompt = random.choice(["1girl, solo", "1boy, muscular", "couple in intimate pose"])

    prompts = []
    for _ in range(num_prompts):
        quality = ", ".join(random.sample(quality_boosters, k=3))  # Selecciona 3 boosters aleatorios
        style = ", ".join(random.sample(styles, k=2))
        nsfw = ", ".join(random.sample(nsfw_tags, k=2))
        light = random.choice(lighting)
        cam = random.choice(camera)
        setting = random.choice(settings)
        prompt = f"{quality}, {input_prompt}, {style}, {nsfw}, {light}, {cam}, {setting}"
        prompts.append(prompt)
    return prompts

# Función para generar imágenes
def generate_images(input_prompt, num_images=3):
    try:
        prompts = generate_prompts(input_prompt, num_images)
        images = []
        for prompt in prompts:
            image = pipe(
                prompt,
                negative_prompt=negative_prompt_base,
                num_inference_steps=30,
                guidance_scale=7.5,
                height=1024,
                width=1024  # Resolución recomendada para Pony XL (múltiplos de 64)
            ).images[0]
            images.append(image)

        # Guardar imágenes
        output_dir = "generated_images"
        os.makedirs(output_dir, exist_ok=True)
        saved_paths = []
        for i, img in enumerate(images):
            path = os.path.join(output_dir, f"image_{i+1}.png")
            img.save(path)
            saved_paths.append(path)

        return images, prompts
    except Exception as e:
        return None, f"Error: {str(e)}"

# Interfaz de Gradio
def create_interface():
    with gr.Blocks(title="Pony XL NSFW Image Generator") as demo:
        gr.Markdown("# Generador de Imágenes Hiperrealistas NSFW con Pony XL")
        gr.Markdown("**Nota:** Este modelo está optimizado para contenido NSFW explícito. Usa prompts en estilo de tags separados por comas para mejores resultados (ej: '1girl, blue eyes, long hair').")
        with gr.Row():
            with gr.Column():
                input_prompt = gr.Textbox(label="Prompt (deja vacío para aleatorio)", placeholder="Ej: 1girl, solo, detailed face")
                num_images = gr.Slider(label="Número de imágenes", minimum=1, maximum=5, step=1, value=3)
                generate_btn = gr.Button("Generar Imágenes")
            with gr.Column():
                output_images = gr.Gallery(label="Imágenes Generadas", show_label=True, elem_id="gallery")
                output_prompts = gr.Textbox(label="Prompts Generados", lines=5, interactive=False)

        generate_btn.click(
            fn=generate_images,
            inputs=[input_prompt, num_images],
            outputs=[output_images, output_prompts]
        )

    return demo

# Lanzar la aplicación
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860)