Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-import spaces
-
 from ominicontrol import generate_image
 
-USE_ZERO_GPU = True
+import spaces
+
+USE_ZERO_GPU = False
 
 css = """
 .inputPanel {
@@ -25,6 +25,15 @@ css = """
 }
 """
 
+header = """
+# 🌍 OminiControl Art
+<div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
+<a href="https://arxiv.org/abs/2411.15098"><img src="https://img.shields.io/badge/ariXv-Paper-A42C25.svg" alt="arXiv"></a>
+<a href="https://github.com/Yuanshi9815/OminiControl"><img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub"></a>
+</div>
+"""
+
+
 
 def style_transfer(image, style):
     return image
@@ -40,7 +49,7 @@ styles = [
 
 def gradio_interface():
     with gr.Blocks(css=css) as demo:
-        gr.Markdown("# 🌍 OminiControl (Image Stylization)")
+        gr.Markdown(header)
         with gr.Row(equal_height=False):
             with gr.Column(variant="panel", elem_classes="inputPanel"):
                 original_image = gr.Image(
@@ -111,6 +120,19 @@ def gradio_interface():
                     label="Output Image",
                 )
 
+        with gr.Row():
+            examples = gr.Examples(
+                examples=[
+                    ["examples/DistractedBoyfriend.webp", styles[0]],
+                    ["examples/steve.webp", styles[0]],
+                    ["examples/oiiai.png", styles[1]],
+                    ["examples/doge.jpg", styles[1]],
+                    ["examples/breakingbad.jpg", styles[2]],
+                    ["examples/PulpFiction.jpg", styles[3]],
+                ],
+                inputs=[original_image, style],
+            )
+
         btn.click(
             fn=infer,
             inputs=[