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, high quality, UHD",),
            ("a beautiful model woman, 20 year aged, Caucasian ethnicity, nurse uniform attire, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a beautiful model woman, Caucasian ethnicity, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a beautiful model woman, Black ethnicity, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),        
            ("a beautiful model woman, Caucasian ethnicity, business attire, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a beautiful model woman, Caucasian ethnicity, casual attire, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a beautiful model woman, golf wear attire, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a handsome model man, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),      
            ("a handsome model man, Caucasian ethnicity, uniform attire, full body visible, model pose, suite wear, direct frontal gaze, blue color background, high quality, UHD",),              
            ("a handsome model man, business attire, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a handsome model man, Caucasian ethnicity, full body visible, model pose, direct frontal gaze, white color background, 16k",),
            ("a handsome model man, Caucasian ethnicity, business attire, Black ethnicity, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),
            ("a handsome model man, casual attire, pilot uniform attire, full body visible, model pose, direct frontal gaze, white color background, high quality, UHD",),              
            ("a handsome model man, Caucasian ethnicity, golf wear attire, full body visible, model pose, direct frontal gaze, gray color background, high quality, UHD",)
        ], 
                    inputs=[prompt], 
                    outputs=[output_gallery])


demo.launch()