Create app.py

app.py

@@ -12,17 +12,13 @@ from PIL import Image, PngImagePlugin
 from datetime import datetime
 from diffusers.models import AutoencoderKL
 from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
-from collections import deque
-import base64
-from io import BytesIO
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 DESCRIPTION = "PonyDiffusion V6 XL"
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n\nRunning on CPU 🥶 This demo does not work on CPU."
-
+    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"
@@ -31,6 +27,7 @@ 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")
+
 MODEL = os.getenv(
     "MODEL",
     "https://huggingface.co/AstraliteHeart/pony-diffusion-v6/blob/main/v6.safetensors",
@@ -41,8 +38,6 @@ torch.backends.cudnn.benchmark = False
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-MAX_HISTORY_SIZE = 10
-image_history = deque(maxlen=MAX_HISTORY_SIZE)
 
 def load_pipeline(model_name):
     vae = AutoencoderKL.from_pretrained(
@@ -54,6 +49,7 @@ def load_pipeline(model_name):
         if MODEL.endswith(".safetensors")
         else StableDiffusionXLPipeline.from_pretrained
     )
+
     pipe = pipeline(
         model_name,
         vae=vae,
@@ -64,9 +60,11 @@ def load_pipeline(model_name):
         use_auth_token=HF_TOKEN,
         variant="fp16",
     )
+
     pipe.to(device)
     return pipe
 
+
 @spaces.GPU
 def generate(
     prompt: str,
@@ -84,16 +82,20 @@ def generate(
     progress=gr.Progress(track_tqdm=True),
 ) -> Image:
     generator = utils.seed_everything(seed)
+
     width, height = utils.aspect_ratio_handler(
-        aspect_ratio_selector, custom_width, custom_height,
+        aspect_ratio_selector,
+        custom_width,
+        custom_height,
     )
+
     width, height = utils.preprocess_image_dimensions(width, height)
+
     backup_scheduler = pipe.scheduler
     pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
 
     if use_upscaler:
         upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
-
     metadata = {
         "prompt": prompt,
         "negative_prompt": negative_prompt,
@@ -115,7 +117,6 @@ def generate(
         }
     else:
         metadata["use_upscaler"] = None
-
     logger.info(json.dumps(metadata, indent=4))
 
     try:
@@ -153,33 +154,12 @@ def generate(
                 output_type="pil",
             ).images
 
-        if images:
-            image = images[0]
-            # Create thumbnail
-            thumbnail = image.copy()
-            thumbnail.thumbnail((256, 256))
-            
-            # Convert thumbnail to base64
-            buffered = BytesIO()
-            thumbnail.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            
-            # Add to history
-            image_history.appendleft({
-                "thumbnail": f"data:image/png;base64,{img_str}",
-                "prompt": prompt,
-                "negative_prompt": negative_prompt,
-                "seed": seed,
-                "width": width,
-                "height": height,
-            })
-
-            if IS_COLAB:
+        if images and IS_COLAB:
+            for image in images:
                 filepath = utils.save_image(image, metadata, OUTPUT_DIR)
                 logger.info(f"Image saved as {filepath} with metadata")
 
-        return image, metadata, list(image_history)
-
+        return images, metadata
     except Exception as e:
         logger.exception(f"An error occurred: {e}")
         raise
@@ -189,6 +169,7 @@ def generate(
         pipe.scheduler = backup_scheduler
         utils.free_memory()
 
+
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)
     logger.info("Loaded on Device!")
@@ -197,32 +178,52 @@ else:
 
 with gr.Blocks(css="style.css") as demo:
     title = gr.HTML(
-        f"""<h1>{DESCRIPTION}</h1>"""
+        f"""<h1><span>{DESCRIPTION}</span></h1>""",
+        elem_id="title",
     )
-
-    with gr.Row():
-        with gr.Column(scale=2):
-            prompt = gr.Textbox(
+    gr.Markdown(
+        f"""Gradio demo for ([Pony Diffusion V6]https://civitai.com/models/257749/pony-diffusion-v6-xl/)""",
+        elem_id="subtitle",
+    )
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
-                max_lines=2,
+                max_lines=5,
                 placeholder="Enter your prompt",
+                container=False,
             )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                show_label=False,
-                max_lines=2,
-                placeholder="Enter a negative prompt",
+            run_button = gr.Button(
+                "Generate", 
+                variant="primary", 
+                scale=0
             )
-
-            with gr.Row():
-                seed = gr.Number(
-                    label="Seed",
-                    value=0,
-                    precision=0,
-                )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
+        result = gr.Gallery(
+            label="Result", 
+            columns=1, 
+            preview=True, 
+            show_label=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=""
+        )
+        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.Row():
                 custom_width = gr.Slider(
                     label="Width",
@@ -238,126 +239,125 @@ with gr.Blocks(css="style.css") as demo:
                     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=1.5,
+                step=0.1,
+                value=1.5,
+                visible=False,
+            )
 
+        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=0, maximum=20, step=0.1, value=7
+                    label="Guidance scale",
+                    minimum=1,
+                    maximum=12,
+                    step=0.1,
+                    value=7.0,
                 )
                 num_inference_steps = gr.Slider(
-                    label="Num Inference Steps",
+                    label="Number of inference steps",
                     minimum=1,
-                    maximum=100,
+                    maximum=50,
                     step=1,
-                    value=30,
-                )
-
-            with gr.Row():
-                sampler = gr.Dropdown(
-                    label="Sampler",
-                    choices=[
-                        "DPM++ 2M SDE Karras",
-                        "DPM++ 2M SDE",
-                        "Euler a",
-                        "Euler",
-                        "DPM++ 2M Karras",
-                        "DPM++ 2M",
-                        "LMS Karras",
-                        "Heun",
-                        "DPM++ SDE Karras",
-                        "DPM++ SDE",
-                        "DPM2 Karras",
-                        "DPM2",
-                        "DPM2 a Karras",
-                        "DPM2 a",
-                        "LMS",
-                        "DDIM",
-                        "PLMS",
-                    ],
-                    value="DPM++ 2M SDE Karras",
-                )
-                aspect_ratio_selector = gr.Dropdown(
-                    label="Aspect Ratio",
-                    choices=[
-                        "1024 x 1024",
-                        "1152 x 896",
-                        "896 x 1152",
-                        "1216 x 832",
-                        "832 x 1216",
-                        "1344 x 768",
-                        "768 x 1344",
-                        "1536 x 640",
-                        "640 x 1536",
-                    ],
-                    value="1024 x 1024",
-                )
-
-            with gr.Row():
-                use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
-                upscaler_strength = gr.Slider(
-                    label="Upscaler Strength",
-                    minimum=0,
-                    maximum=1,
-                    step=0.05,
-                    value=0.55,
-                )
-                upscale_by = gr.Slider(
-                    label="Upscale By",
-                    minimum=1,
-                    maximum=4,
-                    step=0.1,
-                    value=1.5,
+                    value=28,
                 )
-
-        with gr.Column(scale=1):
-            output_image = gr.Image(label="Generated Image")
-            output_text = gr.JSON(label="Generation Info")
-
-    with gr.Row():
-        generate_button = gr.Button("Generate")
-
-    # Add the history component
-    history = gr.HTML(label="Generation History")
-
-    # Update the generate_button click event
-    generate_button.click(
-        generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            custom_width,
-            custom_height,
-            guidance_scale,
-            num_inference_steps,
-            sampler,
-            aspect_ratio_selector,
-            use_upscaler,
-            upscaler_strength,
-            upscale_by,
-        ],
-        outputs=[output_image, output_text, history],
-        concurrency_limit=1
+    with gr.Accordion(label="Generation Parameters", open=False):
+        gr_metadata = gr.JSON(label="Metadata", show_label=False)
+    gr.Examples(
+        examples=config.examples,
+        inputs=prompt,
+        outputs=[result, gr_metadata],
+        fn=lambda *args, **kwargs: generate(*args, use_upscaler=True, **kwargs),
+        cache_examples=CACHE_EXAMPLES,
+    )
+    use_upscaler.change(
+        fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
+        inputs=use_upscaler,
+        outputs=[upscaler_strength, upscale_by],
+        queue=False,
+        api_name=False,
+    )
+    aspect_ratio_selector.change(
+        fn=lambda x: gr.update(visible=x == "Custom"),
+        inputs=aspect_ratio_selector,
+        outputs=custom_resolution,
+        queue=False,
+        api_name=False,
     )
 
-    # Add a function to update the history display
-    def update_history(history_data):
-        html = "<div class='history-container'>"
-        for item in history_data:
-            html += f"""
-            <div class='history-item'>
-                <img src='{item['thumbnail']}' alt='Generated Image'>
-                <div class='history-info'>
-                    <p><strong>Prompt:</strong> {item['prompt']}</p>
-                    <p><strong>Negative Prompt:</strong> {item['negative_prompt']}</p>
-                    <p><strong>Seed:</strong> {item['seed']}</p>
-                    <p><strong>Size:</strong> {item['width']}x{item['height']}</p>
-                </div>
-            </div>
-            """
-        html += "</div>"
-        return html
-
-    # Connect the update_history function to the history component
-    history.change(update_history, inputs=[history], outputs=[history])
-
-demo.launch(debug=True, max_threads=20)
+    inputs = [
+        prompt,
+        negative_prompt,
+        seed,
+        custom_width,
+        custom_height,
+        guidance_scale,
+        num_inference_steps,
+        sampler,
+        aspect_ratio_selector,
+        use_upscaler,
+        upscaler_strength,
+        upscale_by,
+    ]
+
+    prompt.submit(
+        fn=utils.randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
+        fn=generate,
+        inputs=inputs,
+        outputs=result,
+        api_name="run",
+    )
+    negative_prompt.submit(
+        fn=utils.randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
+        fn=generate,
+        inputs=inputs,
+        outputs=result,
+        api_name=False,
+    )
+    run_button.click(
+        fn=utils.randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
+        fn=generate,
+        inputs=inputs,
+        outputs=[result, gr_metadata],
+        api_name=False,
+    )
+demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)