from hidiffusion import apply_hidiffusion, remove_hidiffusion
from diffusers import DiffusionPipeline, DDIMScheduler, AutoencoderKL
import gradio as gr
import torch
import spaces

model = "stabilityai/stable-diffusion-xl-base-1.0"
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
scheduler = DDIMScheduler.from_pretrained(model, subfolder="scheduler")
pipe = DiffusionPipeline.from_pretrained(model, vae=vae, scheduler=scheduler, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")

model_15 = "runwayml/stable-diffusion-v1-5"
scheduler_15 = DDIMScheduler.from_pretrained(model_15, subfolder="scheduler")
pipe_15 = DiffusionPipeline.from_pretrained(model_15, vae=vae, scheduler=scheduler_15, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")

#pipe.enable_model_cpu_offload()
pipe.enable_vae_tiling()

@spaces.GPU
def run_hidiffusion(prompt, negative_prompt="", progress=gr.Progress(track_tqdm=True)):
    apply_hidiffusion(pipe)
    return pipe(prompt, guidance_scale=7.5, height=2048, width=2048, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25).images[0]

@spaces.GPU
def run_hidiffusion_15(prompt, negative_prompt="", progress=gr.Progress(track_tqdm=True)):
    apply_hidiffusion(pipe_15)
    return pipe_15(prompt, guidance_scale=7.5, height=1024, width=1024, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25).images[0]

with gr.Blocks() as demo:
    gr.Markdown("# HiDiffusion Demo")
    with gr.Tab("SDXL in 2048x2048"):
        with gr.Row():
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)",  # 이 부분에 기본값을 설정
                visible=False  # 사용자 인터페이스에서 이 필드를 숨김
            )
        btn = gr.Button("Run")
    with gr.Tab("SD1.5 in 1024x1024"):
        with gr.Row():
            prompt_15 = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)",  # 이 부분에 기본값을 설정
                visible=False  # 사용자 인터페이스에서 이 필드를 숨김
            )
        btn_15 = gr.Button("Run")
    output = gr.Image(label="Result")
    gr.Examples(examples=[
            "a beautiful woman, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, stewardess uniform attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, Caucasian ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, Black ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",        
            "a beautiful woman, Caucasian ethnicity,Business attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, Caucasian ethnicity,casual attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, golf wear attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a handsome man, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",      
            "a handsome man, Caucasian ethnicity,uniform attire, full body visible, model pose, suite wear, direct frontal gaze, blue color background, realritics photo, 16k",              
            "a handsome man, Business attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a handsome man, Caucasian ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a handsome man, Caucasian ethnicity,Business attire, Black ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a handsome man, casual attire, full body visible, model pose, direct frontal gaze, black color background, realritics photo, 16k",              
            "a handsome man, Caucasian ethnicity,golf wear attire, full body visible, model pose, direct frontal gaze, gray color background, realritics photo, 16k"
    ], inputs=[prompt], outputs=[output], fn=run_hidiffusion)
    
    btn.click(fn=run_hidiffusion, inputs=[prompt, negative_prompt], outputs=[output])
    btn_15.click(fn=run_hidiffusion, inputs=[prompt, negative_prompt], outputs=[output])
demo.launch()