import gradio as gr
import numpy as np
import random
import spaces

from diffusers import DiffusionPipeline
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "John6666/wai-nsfw-illustrious-v80-sdxl"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)

if hasattr(pipe, "safety_checker"):
    pipe.safety_checker = None
if hasattr(pipe, "requires_safety_checker"):
    pipe.requires_safety_checker = False
if hasattr(pipe, "enable_attention_slicing"):
    pipe.enable_attention_slicing()
if torch.cuda.is_available() and hasattr(pipe, "enable_xformers_memory_efficient_attention"):
    try:
        pipe.enable_xformers_memory_efficient_attention()
    except Exception:
        pass

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
DEFAULT_NEGATIVE_PROMPT = (
    "worst quality, low quality, lowres, blurry, bad anatomy, bad hands, "
    "extra digits, cropped, watermark, signature, text, jpeg artifacts, 2K censored, 4K censored, "
)


@spaces.GPU
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, seed


examples = [
    "anime heroine, white hair, blue eyes, hooded cloak, glowing orb, full moon, castle background, dynamic pose, masterpiece, best quality, highly detailed",
    "manga swordsman, speed lines, dramatic perspective, night rain, cinematic anime lighting, detailed face, clean line art",
    "anime cafe interior, warm light, slice of life, soft shading, detailed background, expressive eyes, polished illustration",
]

css = """
@import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@500;700&family=Noto+Sans+JP:wght@400;500;700&display=swap');

:root {
    --bg-top: #fff2d9;
    --bg-bottom: #ffd3c7;
    --panel: rgba(255, 250, 244, 0.84);
    --panel-strong: rgba(255, 246, 238, 0.95);
    --ink: #2c1d32;
    --muted: #6c5b67;
    --line: rgba(111, 63, 82, 0.16);
    --accent: #e95f7a;
    --accent-2: #ff9b5c;
    --shadow: 0 26px 70px rgba(124, 66, 83, 0.18);
}

.gradio-container {
    min-height: 100vh;
    font-family: 'Noto Sans JP', sans-serif;
    color: var(--ink);
    background:
        radial-gradient(circle at 12% 14%, rgba(255, 255, 255, 0.95), transparent 24%),
        radial-gradient(circle at 86% 12%, rgba(255, 182, 193, 0.42), transparent 18%),
        linear-gradient(180deg, var(--bg-top) 0%, var(--bg-bottom) 100%);
}

#col-container {
    margin: 28px auto;
    max-width: 900px;
    padding: 28px;
    border: 1px solid var(--line);
    border-radius: 28px;
    background: linear-gradient(180deg, var(--panel-strong), var(--panel));
    box-shadow: var(--shadow);
    backdrop-filter: blur(12px);
}

#title-block h1,
#title-block p {
    margin: 0;
}

#title-block h1 {
    font-family: 'Orbitron', sans-serif;
    font-size: 2.35rem;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    color: #000000;
    opacity: 1;
    -webkit-text-fill-color: #000000;
    text-shadow: none;
}

#title-block p {
    margin-top: 8px;
    font-size: 1rem;
}

#prompt-box textarea,
#neg-box textarea {
    border-radius: 18px;
    background: rgba(255, 255, 255, 0.74);
}

#run-btn {
    min-width: 128px;
    border: 0;
    border-radius: 16px;
    background: linear-gradient(135deg, var(--accent), var(--accent-2));
    box-shadow: 0 14px 30px rgba(233, 95, 122, 0.28);
}

#run-btn:hover {
    filter: brightness(1.04);
}

#result-wrap {
    overflow: hidden;
    border: 1px solid var(--line);
    border-radius: 24px;
    background: linear-gradient(180deg, rgba(255, 255, 255, 0.94), rgba(255, 241, 234, 0.92));
}

.gradio-container .gr-form,
.gradio-container .gr-accordion {
    border-color: var(--line) !important;
}

.gradio-container .gr-accordion summary {
    font-weight: 700;
    color: #7a2f45;
}

.gradio-container button,
.gradio-container input,
.gradio-container textarea {
    font-family: 'Noto Sans JP', sans-serif;
}

@media (max-width: 768px) {
    #col-container {
        margin: 16px;
        padding: 18px;
        border-radius: 22px;
    }

    #title-block h1 {
        font-size: 1.8rem;
    }
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            <div id="title-block">
              <h1>Manga Image Forge</h1>
              <p>Anime and manga inspired image generation with a warmer visual style and cleaner prompt controls.</p>
            </div>
            """
        )

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Describe your anime / manga scene",
                container=False,
                elem_id="prompt-box",
            )

            run_button = gr.Button("Generate", scale=0, variant="primary", elem_id="run-btn")

        result = gr.Image(label="Result", show_label=False, elem_id="result-wrap")

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                value=DEFAULT_NEGATIVE_PROMPT,
                placeholder="Enter a negative prompt",
                visible=True,
                elem_id="neg-box",
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,  # Replace with defaults that work for your model
                )

                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,  # Replace with defaults that work for your model
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=15.0,
                    step=0.1,
                    value=7.0,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=80,
                    step=1,
                    value=28,
                )

        gr.Examples(examples=examples, inputs=[prompt])
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch()