test onediff

.gitignore

@@ -0,0 +1,3 @@
+.vscode
+.DS_Store
+__pycache__

app.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+from app_base import create_demo as create_demo_face
+
+with gr.Blocks(css="style.css") as demo:
+    with gr.Tabs():
+        with gr.Tab(label="Face"):
+            create_demo_face()
+
+demo.launch()

app_onediff.py

@@ -0,0 +1,85 @@
+import spaces
+import gradio as gr
+import time
+import torch
+
+from diffusers import (
+    DDPMScheduler,
+    AutoPipelineForText2Image,
+    AutoencoderTiny,
+)
+
+import oneflow as flow
+from onediff.infer_compiler import oneflow_compile
+
+BASE_MODEL = "stabilityai/sdxl-turbo"
+device = "cuda"
+
+vae = AutoencoderTiny.from_pretrained(
+    'madebyollin/taesdxl',
+    use_safetensors=True,
+    torch_dtype=torch.float16,
+).to('cuda')
+base_pipe = AutoPipelineForText2Image.from_pretrained(
+    BASE_MODEL, 
+    vae=vae,
+    torch_dtype=torch.float16, 
+    variant="fp16", 
+    use_safetensors=True,
+)
+base_pipe.to(device)
+
+base_pipe = base_pipe.to(device, silence_dtype_warnings=True)
+base_pipe.scheduler = DDPMScheduler.from_pretrained(
+    BASE_MODEL,
+    subfolder="scheduler",
+)
+base_pipe.unet = oneflow_compile(base_pipe.unet)
+# base_pipe.vae.decoder = oneflow_compile(base_pipe.vae.decoder)
+
+def create_demo() -> gr.Blocks:
+
+    @spaces.GPU(duration=10)
+    def text_to_image(
+        prompt:str,
+        steps:int,
+    ):
+        run_task_time = 0
+        time_cost_str = ''
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        generated_image = base_pipe(
+            prompt=prompt,
+            num_inference_steps=steps,
+        ).images[0]
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        return generated_image
+
+    def get_time_cost(run_task_time, time_cost_str):
+        now_time = int(time.time()*1000)
+        if run_task_time == 0:
+            time_cost_str = 'start'
+        else:
+            if time_cost_str != '': 
+                time_cost_str += f'-->'
+            time_cost_str += f'{now_time - run_task_time}'
+        run_task_time = now_time
+        return run_task_time, time_cost_str
+
+    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(label="Inference Steps", min=1, max=30, step=1, value=5)
+                g_btn = gr.Button("Generate")
+                
+        with gr.Row():
+            generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+        
+        g_btn.click(
+            fn=text_to_image,
+            inputs=[prompt, steps],
+            outputs=[generated_image],
+        )
+
+    return demo

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+diffusers
+transformers
+accelerate
+spaces