Update app.py

app.py

@@ -1,4 +1,12 @@
+import os
+import gc
+import gradio as gr
+import numpy as np
 import torch
+import json
+import spaces
+import config
+import utils
 import logging
 from PIL import Image, PngImagePlugin
 from datetime import datetime
@@ -12,6 +20,11 @@ DESCRIPTION = "RealVis XL"
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
 IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
+HF_TOKEN = os.getenv("HF_TOKEN")
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
+MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
 
@@ -52,6 +65,10 @@ def load_pipeline(model_name):
 
 @spaces.GPU
 def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    seed: int = 0,
+    custom_width: int = 1024,
     custom_height: int = 1024,
     guidance_scale: float = 7.0,
     num_inference_steps: int = 30,
@@ -131,77 +148,105 @@ with gr.Blocks(css="style.css") as demo:
         elem_id="subtitle",
     )
     gr.DuplicateButton(
-        placeholder="Enter a negative prompt",
-        value=""
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
     )
-    aspect_ratio_selector = gr.Radio(
-        label="Aspect Ratio",
-        choices=config.aspect_ratios,
-        value="1024 x 1024",
-        container=True,
-    )
-    with gr.Group(visible=False) as custom_resolution:
+    with gr.Group():
         with gr.Row():
-            custom_width = gr.Slider(
-                label="Width",
-                minimum=MIN_IMAGE_SIZE,
-                maximum=MAX_IMAGE_SIZE,
-                step=8,
-                value=1024,
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=5,
+                placeholder="Enter your prompt",
+                container=False,
             )
-            custom_height = gr.Slider(
-                label="Height",
-                minimum=MIN_IMAGE_SIZE,
-                maximum=MAX_IMAGE_SIZE,
-                step=8,
-                value=1024,
+            run_button = gr.Button(
+                "Generate", 
+                variant="primary", 
+                scale=0
             )
-    use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
-    with gr.Row() as upscaler_row:
-        upscaler_strength = gr.Slider(
-            label="Strength",
-            minimum=0,
-            maximum=1,
-            step=0.05,
-            value=0.55,
-            visible=False,
+        result = gr.Gallery(
+            label="Result", 
+            columns=1, 
+            preview=True, 
+            show_label=False
         )
-        upscale_by = gr.Slider(
-            label="Upscale by",
-            minimum=1,
-            maximum=1.5,
-            step=0.1,
-            value=1.5,
-            visible=False,
+    with gr.Accordion(label="Advanced Settings", open=False):
+        negative_prompt = gr.Text(
+            label="Negative Prompt",
+            max_lines=5,
+            placeholder="Enter a negative prompt",
+            value=""
         )
-
-    sampler = gr.Dropdown(
-        label="Sampler",
-        choices=config.sampler_list,
-        interactive=True,
-        value="DPM++ 2M SDE Karras",
-    )
-    with gr.Row():
-        seed = gr.Slider(
-            label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
+        aspect_ratio_selector = gr.Radio(
+            label="Aspect Ratio",
+            choices=config.aspect_ratios,
+            value="1024 x 1024",
+            container=True,
         )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-    with gr.Group():
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance scale",
+        with gr.Group(visible=False) as custom_resolution:
+            with gr.Row():
+                custom_width = gr.Slider(
+                    label="Width",
+                    minimum=MIN_IMAGE_SIZE,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=8,
+                    value=1024,
+                )
+                custom_height = gr.Slider(
+                    label="Height",
+                    minimum=MIN_IMAGE_SIZE,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=8,
+                    value=1024,
+                )
+        use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
+        with gr.Row() as upscaler_row:
+            upscaler_strength = gr.Slider(
+                label="Strength",
+                minimum=0,
+                maximum=1,
+                step=0.05,
+                value=0.55,
+                visible=False,
+            )
+            upscale_by = gr.Slider(
+                label="Upscale by",
                 minimum=1,
-                maximum=12,
+                maximum=1.5,
                 step=0.1,
-                value=7.0,
+                value=1.5,
+                visible=False,
             )
-            num_inference_steps = gr.Slider(
-                label="Number of inference steps",
-                minimum=1,
-                maximum=100,
-                step=1,
-                value=28,
+
+        sampler = gr.Dropdown(
+            label="Sampler",
+            choices=config.sampler_list,
+            interactive=True,
+            value="DPM++ 2M SDE Karras",
+        )
+        with gr.Row():
+            seed = gr.Slider(
+                label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
             )
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        with gr.Group():
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=1,
+                    maximum=12,
+                    step=0.1,
+                    value=7.0,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=28,
+                )
     with gr.Accordion(label="Generation Parameters", open=False):
         gr_metadata = gr.JSON(label="Metadata", show_label=False)
     gr.Examples(
@@ -229,6 +274,8 @@ with gr.Blocks(css="style.css") as demo:
     inputs = [
         prompt,
         negative_prompt,
+        seed,
+        custom_width,
         custom_height,
         guidance_scale,
         num_inference_steps,
@@ -275,5 +322,4 @@ with gr.Blocks(css="style.css") as demo:
         outputs=[result, gr_metadata],
         api_name=False,
     )
-
-demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB, show_error=True)
+demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB, show_error=True)