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 = "stabilityai/stable-diffusion-3-medium-diffusers"

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

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

MAX_IMAGE_SIZE = 512

examples = [
    "A border collie lying in some Fall leaves as the forest trees change colors",
    "A cyberpunk cityscape with neon lights and flying cars at night",
    "A fantasy castle surrounded by mountains under a pink and purple sunset",
]


@spaces.GPU(duration=30)
def infer(
    prompt,
    width=512,
    height=512,
    num_inference_steps=20,
    guidance_scale=7.5,
):
    # Генерация изображения
    image = pipe(
        prompt=prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        width=width,
        height=height,
    ).images[0]
    return image


class CustomTheme(gr.themes.Base):
    def __init__(self):
        super().__init__()
        self.primary_hue = "#17181B"
        self.background_fill_primary = "#17181B"
        self.background_fill_secondary = "#17181B"
        self.background_fill_tertiary = "#17181B"
        self.text_color_primary = "#FFFFFF"
        self.text_color_secondary = "#FFFFFF"
        self.text_color_tertiary = "#FFFFFF"
        self.input_background_fill = "#17181B"
        self.input_text_color = "#FFFFFF"


css = """
/* Скрываем нижний колонтитул */
footer {
    visibility: hidden;
    height: 0;
    margin: 0;
    padding: 0;
    overflow: hidden;
}
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@400;500;700&display=swap');

/* Применяем шрифты */
body, input, button, textarea, select, .gr-button {
    font-family: 'Poppins', sans-serif;
    background-color: #17181B !important;
    color: #FFFFFF;
}

/* Настройки заголовков */
h1, h2, h3, h4, h5, h6 {
    font-family: 'Poppins', sans-serif;
    font-weight: 700;
    color: #FFFFFF;
}

/* Стиль для текстовых полей и кнопок */
input[type="text"], textarea {
    background-color: #17181B !important;
    color: #FFFFFF;
    border: 1px solid #FFFFFF;
}

/* Цвет кнопки Generate */
.generate-button {
    background-color: #5271FF !important;
    color: #FFFFFF !important;
    border: none;
    font-weight: bold;
}

.generate-button:hover {
    background-color: #405BBF !important; /* Цвет при наведении */
}

/* Выделяем текст для Prompt */
.prompt-text {
    font-weight: bold;
    color: #FFFFFF;
}
"""

with gr.Blocks(theme=CustomTheme(), css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("**Prompt**", elem_classes="prompt-text")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )

            run_button = gr.Button(
                "Generate",
                scale=0,
                variant="primary",
                elem_classes="generate-button",
            )

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

        gr.Examples(
            examples=examples,
            inputs=[prompt],
            outputs=[result],
            fn=infer,
            cache_examples=True,
            cache_mode="lazy",
        )

        run_button.click(
            fn=infer,
            inputs=[prompt],
            outputs=[result],
        )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_api=False,
    )