import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
import transformers
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
transformers.logging.set_verbosity_error()

# 모델 경로 목록
model_paths = {
    "CompVis/stable-diffusion-v1-4": "CompVis/stable-diffusion-v1-4",
    "runwayml/stable-diffusion-v1-5": "runwayml/stable-diffusion-v1-5",
    "stabilityai/stable-diffusion-2-1": "stabilityai/stable-diffusion-2-1"
}

# 모델을 로드합니다.
pipes = {}
for name, path in model_paths.items():
    pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
    pipes[name] = pipe

def generate_image(prompt, negative_prompt, model_name, steps, sampler, cfg_scale, width, height, seed):
    pipe = pipes[model_name]
    if seed == -1:
        seed = None
    generator = torch.manual_seed(seed) if seed is not None else None

    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        width=width,
        height=height,
        generator=generator,
    )

    image = output.images[0]

    return image


css = """
#generate {
    height: 100%;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Row():
        model_name = gr.Dropdown(
            label="Model (모델 선택)",
            choices=list(model_paths.keys()),
            value="runwayml/stable-diffusion-v1-5",
            interactive=True
        )

    with gr.Tab("Text-to-Image (텍스트에서 이미지 생성)"):
        with gr.Row():
            with gr.Column(scale=6, min_width=600):
                prompt = gr.Textbox(
                    label="Prompt (프롬프트)",
                    value="a high-resolution photograph of a samoyed cub wearing a shiny silver astronaut suit",
                    lines=3
                )
                negative_prompt = gr.Textbox(
                    label="Negative Prompt (네거티브 프롬프트)",
                    value="low resolution, blurry, cartoon, deformed, ugly, disfigured, malformed, bad anatomy, unrealistic",
                    lines=3
                )
            with gr.Column():
                generate_button = gr.Button("Generate (생성)", variant='primary')

        with gr.Row():
            with gr.Column(scale=3):
                with gr.Tab("Generation Settings (생성 설정)"):
                    with gr.Row():
                        sampler = gr.Dropdown(
                            label="Sampler (샘플러)",
                            choices=["ddim", "ddpm", "pndm"],
                            value="ddpm"
                        )

                    with gr.Row():
                        steps = gr.Slider(
                            label="Steps (샘플링 단계 수)",
                            minimum=1,
                            maximum=50,
                            value=20,
                            step=1
                        )

                    with gr.Row():
                        width = gr.Slider(
                            label="Width (너비, 픽셀)",
                            minimum=512,
                            maximum=1024,
                            value=512,
                            step=64
                        )
                        height = gr.Slider(
                            label="Height (높이, 픽셀)",
                            minimum=512,
                            maximum=1024,
                            value=512,
                            step=64
                        )

                    cfg_scale = gr.Slider(
                        label="CFG Scale (CFG 스케일, 제어 강도)",
                        minimum=1,
                        maximum=20,
                        value=7,
                        step=1
                    )
                    seed = gr.Number(
                        label="Seed (시드, 고유 번호)",
                        value=-1
                    )

            with gr.Column(scale=2):
                output_image = gr.Image(
                    label="Output Image (생성된 이미지)"
                )

        generate_button.click(
            fn=generate_image,
            inputs=[
                prompt,
                negative_prompt,
                model_name,
                steps,
                sampler,
                cfg_scale,
                width,
                height,
                seed
            ],
            outputs=output_image
        )

demo.launch()