History v1

app.py

@@ -12,13 +12,17 @@ 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<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
+    DESCRIPTION += "\n\nRunning on CPU 🥶 This demo does not work on CPU."
+
 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"
@@ -27,7 +31,6 @@ 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",
@@ -38,6 +41,8 @@ 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(
@@ -49,7 +54,6 @@ def load_pipeline(model_name):
         if MODEL.endswith(".safetensors")
         else StableDiffusionXLPipeline.from_pretrained
     )
-
     pipe = pipeline(
         model_name,
         vae=vae,
@@ -60,11 +64,9 @@ def load_pipeline(model_name):
         use_auth_token=HF_TOKEN,
         variant="fp16",
     )
-
     pipe.to(device)
     return pipe
 
-
 @spaces.GPU
 def generate(
     prompt: str,
@@ -82,20 +84,16 @@ 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,
@@ -117,6 +115,7 @@ def generate(
         }
     else:
         metadata["use_upscaler"] = None
+
     logger.info(json.dumps(metadata, indent=4))
 
     try:
@@ -154,12 +153,34 @@ def generate(
                 output_type="pil",
             ).images
 
+        if images:
+            for image in images:
+                # 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 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 images, metadata
+        return images, metadata, list(image_history)
+
     except Exception as e:
         logger.exception(f"An error occurred: {e}")
         raise
@@ -169,7 +190,6 @@ def generate(
         pipe.scheduler = backup_scheduler
         utils.free_memory()
 
-
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)
     logger.info("Loaded on Device!")
@@ -178,52 +198,32 @@ else:
 
 with gr.Blocks(css="style.css") as demo:
     title = gr.HTML(
-        f"""<h1><span>{DESCRIPTION}</span></h1>""",
-        elem_id="title",
-    )
-    gr.Markdown(
-        f"""Gradio demo for ([Pony Diffusion V6]https://civitai.com/models/257749/pony-diffusion-v6-xl/)""",
-        elem_id="subtitle",
+        f"""<h1>{DESCRIPTION}</h1>"""
     )
-    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(
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt = gr.Textbox(
                 label="Prompt",
                 show_label=False,
-                max_lines=5,
+                max_lines=2,
                 placeholder="Enter your prompt",
-                container=False,
             )
-            run_button = gr.Button(
-                "Generate", 
-                variant="primary", 
-                scale=0
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                show_label=False,
+                max_lines=2,
+                placeholder="Enter a negative prompt",
             )
-        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():
+                seed = gr.Number(
+                    label="Seed",
+                    value=0,
+                    precision=0,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
             with gr.Row():
                 custom_width = gr.Slider(
                     label="Width",
@@ -239,125 +239,126 @@ 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=1,
-                    maximum=12,
-                    step=0.1,
-                    value=7.0,
+                    label="Guidance Scale", minimum=0, maximum=20, step=0.1, value=7
                 )
                 num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
+                    label="Num Inference Steps",
                     minimum=1,
-                    maximum=50,
+                    maximum=100,
                     step=1,
-                    value=28,
+                    value=30,
                 )
-    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,
-    )
 
-    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,
+            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,
+                )
+
+        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],
     )
-demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)
+
+    # 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.queue(concurrency_count=1, max_size=20)
+demo.launch(debug=True)