app.py

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, full body visible, model pose, direct frontal gaze, red 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, Business attire, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a beautiful woman, 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, full body visible, model pose, suite wear, direct frontal gaze, blue 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, full body visible, model pose, direct frontal gaze, white color background, realritics photo, 16k",
            "a handsome man, Business attire, Black ethnicity, 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, black color background, realritics photo, 16k",              
            "a handsome man, 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()