| import spaces |
| import gradio as gr |
| import torch |
| import os |
|
|
| from diffusers import ( |
| DDPMScheduler, |
| StableDiffusionXLImg2ImgPipeline, |
| AutoencoderKL, |
| ) |
|
|
| from diffusers.utils import load_image |
|
|
| os.system("pip install torch_tensorrt==2.4.0") |
|
|
| BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0" |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"--------->Device: {device}") |
|
|
| vae = AutoencoderKL.from_pretrained( |
| "madebyollin/sdxl-vae-fp16-fix", |
| torch_dtype=torch.float16, |
| ) |
|
|
| base_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained( |
| BASE_MODEL, |
| vae=vae, |
| torch_dtype=torch.float16, |
| variant="fp16", |
| use_safetensors=True, |
| ) |
| base_pipe = base_pipe.to(device, silence_dtype_warnings=True) |
| base_pipe.scheduler = DDPMScheduler.from_pretrained( |
| BASE_MODEL, |
| subfolder="scheduler", |
| ) |
|
|
| backend = "torch_tensorrt" |
| import torch_tensorrt |
| print('Compiling model...') |
| compiledModel = torch.compile( |
| base_pipe.unet, |
| backend=backend, |
| options={ |
| "truncate_long_and_double": True, |
| "enabled_precisions": {torch.float32, torch.float16}, |
| }, |
| dynamic=False, |
| ) |
|
|
| base_pipe.unet = compiledModel |
|
|
| init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img.png") |
| generated_image = base_pipe( |
| init_image, |
| num_inference_steps=5, |
| ).images[0] |
|
|
| generated_image.save("/tmp/gradio/generated_image.png") |
|
|
|
|
| def create_demo() -> gr.Blocks: |
|
|
| @spaces.GPU(duration=30) |
| def text_to_image( |
| prompt:str, |
| steps:int, |
| ): |
| import torch_tensorrt |
| print('Compiling model...') |
| compiledModel = torch.compile( |
| base_pipe.unet, |
| backend=backend, |
| options={ |
| "truncate_long_and_double": True, |
| "enabled_precisions": {torch.float32, torch.float16}, |
| }, |
| dynamic=False, |
| ) |
| print('Model compiled!') |
|
|
| print('Saving compiled model...') |
| torch_tensorrt.save(compiledModel, "compiled_pipe.ep") |
| print('Compiled model saved!') |
|
|
| with gr.Blocks() as demo: |
| with gr.Row(): |
| with gr.Column(): |
| prompt = gr.Textbox(label="Prompt", placeholder="Write a prompt here", lines=2, value="A beautiful sunset over the city") |
| with gr.Column(): |
| steps = gr.Slider(minimum=1, maximum=100, value=5, step=1, label="Num Steps") |
| g_btn = gr.Button("Generate") |
| |
| with gr.Row(): |
| with gr.Column(): |
| generated_image = gr.Image(label="Generated Image", type="pil", interactive=False) |
| with gr.Column(): |
| time_cost = gr.Textbox(label="Time Cost", lines=1, interactive=False) |
| |
| g_btn.click( |
| fn=text_to_image, |
| inputs=[prompt, steps], |
| |
| outputs=[], |
| ) |
|
|
| return demo |
|
|
