import torch
import spaces
import gradio as gr
from diffusers import StableDiffusionXLPipeline, AutoencoderKL

# Ссылки на файлы из репозиториев
BASE_MODEL_URL = "https://huggingface.co/LyliaEngine/Pony_Diffusion_V6_XL/resolve/main/ponyDiffusionV6XL_v6StartWithThisOne.safetensors"
VAE_MODEL_URL = "https://huggingface.co/LyliaEngine/Pony_Diffusion_V6_XL/resolve/main/sdxl_vae.safetensors"
LORA_MODEL_URL = "https://huggingface.co/nncyberpunk/LoRA_SDXL1.0_SummertimeSagaXL_Pony/resolve/main/LoRA_SDXL1.0_SummertimeSagaXL_Pony.safetensors"

# 1. Загружаем VAE напрямую
# Это гарантирует правильные цвета как в оригинальной Pony V6
vae = AutoencoderKL.from_single_file(VAE_MODEL_URL, torch_dtype=torch.float16)

# 2. Инициализируем пайплайн с этим VAE
pipe = StableDiffusionXLPipeline.from_single_file(
    BASE_MODEL_URL, 
    vae=vae,
    torch_dtype=torch.float16,
    use_safetensors=True
)

# 3. Накладываем LoRA
pipe.load_lora_weights(LORA_MODEL_URL)
pipe.to("cuda")

@spaces.GPU
def generate(prompt, negative_prompt, lora_scale, steps, guidance):
    # Добавляем триггеры для Pony и стиля Summertime Saga
    # Важно: 'summertime saga style' активирует LoRA
    styled_prompt = f"score_9, score_8, score_7, rating_explicit, summertime saga style, {prompt}"
    
    image = pipe(
        prompt=styled_prompt,
        negative_prompt=negative_prompt,
        cross_attention_kwargs={"scale": lora_scale},
        num_inference_steps=steps,
        guidance_scale=guidance,
    ).images[0]
    return image

# Интерфейс (Gradio)
with gr.Blocks() as demo:
    gr.Markdown("### Summertime Saga XL + Pony V6 (Custom VAE)")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Промпт", placeholder="jenny, bedroom, looking at viewer...")
            neg_prompt = gr.Textbox(
                label="Negative Prompt", 
                value="score_6, score_5, score_4, low quality, bad anatomy, 3d, render, photo"
            )
            
            with gr.Row():
                lora_scale = gr.Slider(0, 1.2, value=0.85, label="LoRA Strength")
                steps = gr.Slider(10, 50, value=25, label="Steps")
            
            guidance = gr.Slider(1, 10, value=7.0, label="CFG Scale")
            btn = gr.Button("Generate", variant="primary")
            
        with gr.Column():
            output = gr.Image(label="Результат")

    btn.click(
        fn=generate, 
        inputs=[prompt, neg_prompt, lora_scale, steps, guidance], 
        outputs=output
    )

demo.launch()