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_vae_tiling()

@spaces.GPU
def run_hidiffusion(prompt, negative_prompt="", progress=gr.Progress(track_tqdm=True)):
    apply_hidiffusion(pipe)
    images = []
    for _ in range(2):  # 4개의 이미지 생성
        result = pipe(prompt, guidance_scale=7.5, height=2048, width=2048, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25)
        images.append(result.images[0])
    return images

 
with gr.Blocks() as demo:
    gr.Markdown("# AI Human Model Generator")
     
    with gr.Tab("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")
        output_gallery = gr.Gallery(label="Results")
        btn.click(fn=run_hidiffusion, inputs=[prompt, negative_prompt], outputs=[output_gallery])

        gr.Examples(examples=[
            "a beautiful model woman, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a beautiful model woman, 20 year aged, Caucasian ethnicity, nurse uniform attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo,High Quaility, UHD",
            "a beautiful model woman, Caucasian ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a beautiful model woman, Black ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",        
            "a beautiful model woman, Caucasian ethnicity,Business attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a beautiful model woman, Caucasian ethnicity,casual attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a beautiful model woman, golf wear attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a handsome model man, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",      
            "a handsome model man, Caucasian ethnicity,uniform attire, full body visible, model pose, suite wear, direct frontal gaze, blue color background, realritics photo, High Quaility, UHD",              
            "a handsome model man, Business attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a handsome model man, Caucasian ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a handsome model man, Caucasian ethnicity,Business attire, Black ethnicity, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",
            "a handsome model man, casual attire, pilot uniform attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, High Quaility, UHD",              
            "a handsome model man, Caucasian ethnicity,golf wear attire, full body visible, model pose, direct frontal gaze, gray color background, realritics photo, High Quaility, UHD"
        ], 
                    inputs=[prompt], 
                    outputs=[output_gallery])

demo.launch()