add app_tensorrt

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from app_onediff import create_demo as create_demo_face
+from app_tensorrt import create_demo as create_demo_face
 
 with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():

app_tensorrt.py

@@ -0,0 +1,85 @@
+import torch
+import torch_tensorrt
+
+from diffusers import (
+    DDPMScheduler,
+    StableDiffusionXLImg2ImgPipeline,
+    AutoencoderKL,
+)
+
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+device = "cuda"
+
+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"
+
+
+
+# print('Loading compiled model...')
+# loadedModel = torch_tensorrt.load("compiled_pipe.ep").module()
+# print('Compiled model loaded!')
+
+
+
+def create_demo() -> gr.Blocks:
+
+    @spaces.GPU(duration=30)
+    def text_to_image(
+        prompt:str,
+        steps:int,
+    ):
+        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=[generated_image, time_cost],
+            outputs=[],
+        )
+
+    return demo

requirements.txt

@@ -1,13 +1,8 @@
 gradio
-torch
-torchvision
+torch==2.5.0
+torch_tensorrt==2.5.0
+torchvision==0.20.0
 diffusers
 transformers
 accelerate
-spaces
-git+https://github.com/XPixelGroup/BasicSR@master
-gfpgan
-facexlib
-realesrgan
-triton
-xformers
+spaces