import spaces  # <-- DAPAT ITO ANG PINAKAUNANG LINYA BAGO ANG LAHAT!

import gc
import torch
import gradio as gr
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

# =========================================================================
# LAHAT NG 8 CHECKPOINT MODELS MULA SA HUGGING FACE
# =========================================================================
MODELS = {
    "SpatterXL v1.7 (Illustrious x Pony)": {
        "type": "pretrained",
        "path": "John6666/spatterxl-illustriousxl-ponyxl-v15-sdxl"
    },
    "Solmeleon-XL (NoobAI x Illustrious)": {
        "type": "pretrained",
        "path": "John6666/solmeleon-illustriousxl-v10-sdxl"
    },
    "HoneyMix PonyXL (High Contrast Anime)": {
        "type": "pretrained",
        "path": "John6666/honey-mix-xl-high-contrast-anime-checkpoint-v1-sdxl"
    },
    "Pony Diffusion V6 XL (Waifu & NSFW)": {
        "type": "single_file",
        "path": "https://huggingface.co/LyliaEngine/Pony_Diffusion_V6_XL/resolve/main/ponyDiffusionV6XL_v6StartWithThisOne.safetensors"
    },
    "Illustrious-XL v0.1 (Official Base)": {
        "type": "single_file",
        "path": "https://huggingface.co/OnomaAIResearch/Illustrious-xl-early-release-v0/resolve/main/Illustrious-XL-v0.1.safetensors"
    },
    "Hoseki LustrousMix (Vibrant Anime)": {
        "type": "single_file",
        "path": "https://huggingface.co/LyliaEngine/hosekiLustrousmixAnima_animaV10/resolve/main/hosekiLustrousmixAnima_animaV10.safetensors"
    },
    "Initium-XL (NoobAI Epsilon 1.1)": {
        "type": "single_file",
        "path": "https://huggingface.co/John6666/initium-finetune-noobaiepsilon11-v10-sdxl/resolve/main/initium-finetune-noobaiepsilon11-v10-sdxl.safetensors"
    },
    "Tryptophan-666 (Illustrious x Pony XL)": {
        "type": "single_file",
        "path": "https://huggingface.co/John6666/tryptophan-666-illustrious-xl-pony-compatable-v1666-sdxl/resolve/main/tryptophan-666-illustrious-xl-pony-compatable-v1666-sdxl.safetensors"
    }
}

current_model_name = None
pipe = None

def load_pipeline(model_choice):
    global pipe, current_model_name
    if pipe is None or current_model_name != model_choice:
        if pipe is not None:
            del pipe
            gc.collect()
            torch.cuda.empty_cache()

        cfg = MODELS[model_choice]
        print(f"Loading checkpoint: {model_choice}...")
        
        if cfg["type"] == "pretrained":
            pipe = StableDiffusionXLPipeline.from_pretrained(
                cfg["path"],
                torch_dtype=torch.float16,
                use_safetensors=True
            )
        else:
            pipe = StableDiffusionXLPipeline.from_single_file(
                cfg["path"],
                torch_dtype=torch.float16,
                use_safetensors=True
            )
        
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        current_model_name = model_choice
    return pipe

# ZeroGPU: Dynamic Nvidia A100 GPU Allocation
@spaces.GPU(duration=45)
def generate_image(prompt, negative_prompt, model_choice, width, height, steps, cfg_scale, seed):
    pipeline = load_pipeline(model_choice)
    pipeline.to("cuda")
    
    generator = None
    if seed != -1:
        generator = torch.Generator(device="cuda").manual_seed(int(seed))
        
    is_pony_family = any(k in model_choice for k in ["Pony", "Spatter", "Honey", "Tryptophan", "Initium"])
    if is_pony_family and "score_9" not in prompt:
        prompt = "score_9, score_8_up, score_7_up, source_anime, " + prompt + ", masterpiece, best quality"
        negative_prompt = "score_4, score_5, score_6, ugly, deformed, blurry, bad anatomy, " + negative_prompt
    elif "masterpiece" not in prompt:
        prompt = "masterpiece, best quality, very aesthetic, " + prompt
        negative_prompt = "ugly, deformed, blurry, bad hands, low quality, " + negative_prompt

    image = pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=int(width),
        height=int(height),
        num_inference_steps=int(steps),
        guidance_scale=float(cfg_scale),
        generator=generator
    ).images[0]
    
    return image

custom_css = """
body, .gradio-container { background-color: #141414 !important; color: #ffffff !important; }
button.primary { background: linear-gradient(90deg, #E50914 0%, #b80710 100%) !important; border: none !important; color: white !important; font-weight: bold !important; }
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
    gr.Markdown("""
    # 👑 QUEENDOM STUDIO — ZeroGPU Backend (8 Models)
    High-Performance SDXL, Pony, and Illustrious Generator powered by free dynamic Nvidia A100.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            model_dd = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="SpatterXL v1.7 (Illustrious x Pony)",
                label="Select Checkpoint Model"
            )
            prompt_input = gr.Textbox(label="Positive Prompt", lines=3, value="1girl, solo, silver hair, golden eyes, beautiful anime waifu")
            neg_prompt_input = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, deformed, blurry, bad anatomy")
            
            with gr.Row():
                width_slider = gr.Slider(512, 1216, value=832, step=64, label="Width")
                height_slider = gr.Slider(512, 1216, value=1216, step=64, label="Height")
            
            with gr.Row():
                steps_slider = gr.Slider(15, 40, value=26, step=1, label="Steps")
                cfg_slider = gr.Slider(3.0, 10.0, value=6.0, step=0.5, label="CFG Scale")
                
            seed_input = gr.Number(value=-1, label="Seed (-1 for Random)")
            generate_btn = gr.Button("✨ GENERATE ARTWORK", variant="primary", elem_classes=["primary"])
            
        with gr.Column(scale=1):
            output_img = gr.Image(label="Rendered Artwork", type="pil")
            
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, neg_prompt_input, model_dd, width_slider, height_slider, steps_slider, cfg_slider, seed_input],
        outputs=output_img
    )

demo.queue().launch()