v3p4

app.py

@@ -1,6 +1,5 @@
 import os
 import gc
-import random
 import gradio as gr
 import numpy as np
 import torch
@@ -14,82 +13,24 @@ from datetime import datetime
 from diffusers.models import AutoencoderKL
 from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
 
-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>"
-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")
-
-MODEL = os.getenv(
-    "MODEL",
-    "https://huggingface.co/AstraliteHeart/pony-diffusion-v6/blob/main/v6.safetensors",
-)
-
-torch.backends.cudnn.deterministic = True
-torch.backends.cudnn.benchmark = False
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-def load_pipeline(model_name):
-    vae = AutoencoderKL.from_pretrained(
-        "madebyollin/sdxl-vae-fp16-fix",
-        torch_dtype=torch.float16,
-    )
-    pipeline = (
-        StableDiffusionXLPipeline.from_single_file
-        if MODEL.endswith(".safetensors")
-        else StableDiffusionXLPipeline.from_pretrained
-    )
-
-    pipe = pipeline(
-        model_name,
-        vae=vae,
-        torch_dtype=torch.float16,
-        custom_pipeline="lpw_stable_diffusion_xl",
-        use_safetensors=True,
-        add_watermarker=False,
-        use_auth_token=HF_TOKEN,
-        variant="fp16",
-    )
-
-    pipe.to(device)
-    return pipe
+# ... (keep the existing imports and configurations)
 
+# Add a new function to parse and validate JSON input
 def parse_json_parameters(json_str):
     try:
         params = json.loads(json_str)
+        required_keys = ['prompt', 'negative_prompt', 'seed', 'width', 'height', 'guidance_scale', 'num_inference_steps', 'sampler']
+        for key in required_keys:
+            if key not in params:
+                raise ValueError(f"Missing required key: {key}")
         return params
     except json.JSONDecodeError:
-        return None
-
-def apply_json_parameters(json_str):
-    params = parse_json_parameters(json_str)
-    if params:
-        return (
-            params.get("prompt", ""),
-            params.get("negative_prompt", ""),
-            params.get("seed", 0),
-            params.get("width", 1024),
-            params.get("height", 1024),
-            params.get("guidance_scale", 7.0),
-            params.get("num_inference_steps", 30),
-            params.get("sampler", "DPM++ 2M SDE Karras"),
-            params.get("aspect_ratio", "1024 x 1024"),
-            params.get("use_upscaler", False),
-            params.get("upscaler_strength", 0.55),
-            params.get("upscale_by", 1.5),
-        )
-    return [gr.update()] * 12
+        raise ValueError("Invalid JSON format")
+    except Exception as e:
+        raise ValueError(f"Error parsing JSON: {str(e)}")
 
+# Modify the generate function to accept JSON parameters
+@spaces.GPU
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -103,8 +44,23 @@ def generate(
     use_upscaler: bool = False,
     upscaler_strength: float = 0.55,
     upscale_by: float = 1.5,
+    json_params: str = "",
     progress=gr.Progress(track_tqdm=True),
 ) -> Image:
+    if json_params:
+        try:
+            params = parse_json_parameters(json_params)
+            prompt = params['prompt']
+            negative_prompt = params['negative_prompt']
+            seed = params['seed']
+            custom_width = params['width']
+            custom_height = params['height']
+            guidance_scale = params['guidance_scale']
+            num_inference_steps = params['num_inference_steps']
+            sampler = params['sampler']
+        except ValueError as e:
+            raise gr.Error(str(e))
+
     generator = utils.seed_everything(seed)
 
     width, height = utils.aspect_ratio_handler(
@@ -183,7 +139,7 @@ def generate(
                 filepath = utils.save_image(image, metadata, OUTPUT_DIR)
                 logger.info(f"Image saved as {filepath} with metadata")
 
-        return images, json.dumps(metadata)  # Return metadata as a JSON string
+        return images, metadata
     except Exception as e:
         logger.exception(f"An error occurred: {e}")
         raise
@@ -193,21 +149,30 @@ def generate(
         pipe.scheduler = backup_scheduler
         utils.free_memory()
 
-def get_random_prompt():
-    anime_characters = [
-        "Naruto Uzumaki", "Monkey D. Luffy", "Goku", "Eren Yeager", "Light Yagami",
-        "Lelouch Lamperouge", "Edward Elric", "Levi Ackerman", "Spike Spiegel",
-        "Sakura Haruno", "Mikasa Ackerman", "Asuka Langley Soryu", "Rem", "Megumin",
-        "Violet Evergarden"
-    ]
-    styles = ["pixel art", "stylized anime", "digital art", "watercolor", "sketch"]
-    scores = ["score_9", "score_8_up", "score_7_up"]
-    
-    character = random.choice(anime_characters)
-    style = random.choice(styles)
-    score = ", ".join(random.sample(scores, k=3))
-    
-    return f"{score}, {character}, {style}, show accurate"
+# Initialize an empty list to store the generation history
+generation_history = []
+
+# Function to update the history dropdown
+def update_history_dropdown():
+    return gr.Dropdown.update(choices=[f"{item['prompt']} ({item['timestamp']})" for item in generation_history])
+
+# Modify the generate function to add results to the history
+def generate_and_update_history(*args, **kwargs):
+    images, metadata = generate(*args, **kwargs)
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    generation_history.insert(0, {"prompt": metadata["prompt"], "timestamp": timestamp, "image": images[0], "metadata": metadata})
+    if len(generation_history) > 10:  # Limit history to 10 items
+        generation_history.pop()
+    return images, metadata, update_history_dropdown()
+
+# Function to display selected history item
+def display_history_item(selected_item):
+    if not selected_item:
+        return None, None
+    for item in generation_history:
+        if f"{item['prompt']} ({item['timestamp']})" == selected_item:
+            return item['image'], json.dumps(item['metadata'], indent=2)
+    return None, None
 
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)
@@ -215,59 +180,7 @@ if torch.cuda.is_available():
 else:
     pipe = None
 
-# Define the JavaScript code as a string
-js_code = """
-<script>
-document.addEventListener('DOMContentLoaded', (event) => {
-  const historyDropdown = document.getElementById('history-dropdown');
-  const resultGallery = document.querySelector('.gallery');
-
-  if (historyDropdown && resultGallery) {
-    const observer = new MutationObserver((mutations) => {
-      mutations.forEach((mutation) => {
-        if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {
-          const newImage = mutation.addedNodes[0];
-          if (newImage.tagName === 'IMG') {
-            updateHistory(newImage.src);
-          }
-        }
-      });
-    });
-
-    observer.observe(resultGallery, { childList: true });
-
-    function updateHistory(imageSrc) {
-      const prompt = document.querySelector('#prompt textarea').value;
-      const option = document.createElement('option');
-      option.value = prompt;
-      option.textContent = prompt;
-      option.setAttribute('data-image', imageSrc);
-
-      historyDropdown.insertBefore(option, historyDropdown.firstChild);
-
-      if (historyDropdown.children.length > 10) {
-        historyDropdown.removeChild(historyDropdown.lastChild);
-      }
-    }
-
-    historyDropdown.addEventListener('change', (event) => {
-      const selectedOption = event.target.selectedOptions[0];
-      const imageSrc = selectedOption.getAttribute('data-image');
-      if (imageSrc) {
-        const img = document.createElement('img');
-        img.src = imageSrc;
-        resultGallery.innerHTML = '';
-        resultGallery.appendChild(img);
-      }
-    });
-  }
-});
-</script>
-"""
-
 with gr.Blocks(css="style.css") as demo:
-    gr.HTML(js_code)  # Add the JavaScript code to the interface
-    
     title = gr.HTML(
         f"""<h1><span>{DESCRIPTION}</span></h1>""",
         elem_id="title",
@@ -376,34 +289,24 @@ with gr.Blocks(css="style.css") as demo:
                     step=1,
                     value=28,
                 )
-    with gr.Accordion(label="JSON Parameters", open=False):
-        json_input = gr.TextArea(label="Input JSON parameters")
-        apply_json_button = gr.Button("Apply JSON Parameters")
-
-    with gr.Row():
-        clear_button = gr.Button("Clear All")
-        random_prompt_button = gr.Button("Random Prompt")
-
-    history = gr.State([])  # Add a state component to store history
-    history_dropdown = gr.Dropdown(label="Generation History", choices=[], interactive=True, elem_id="history-dropdown")
-    
     with gr.Accordion(label="Generation Parameters", open=False):
         gr_metadata = gr.JSON(label="Metadata", show_label=False)
-    
-    def update_history(images, metadata, history):
-        if images:
-            new_entry = {"prompt": json.loads(metadata)["prompt"], "image": images[0]}
-            history = [new_entry] + history[:9]  # Keep only the last 10 entries
-        return gr.update(choices=[h["prompt"] for h in history]), history
+        json_input = gr.TextArea(label="Edit/Paste JSON Parameters", placeholder="Paste or edit JSON parameters here")
+        generate_from_json = gr.Button("Generate from JSON")
+
+    # Add history dropdown
+    history_dropdown = gr.Dropdown(label="Generation History", choices=[], interactive=True)
+    history_image = gr.Image(label="Selected Image", interactive=False)
+    history_metadata = gr.JSON(label="Selected 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),
+        fn=lambda *args, **kwargs: generate_and_update_history(*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,
@@ -432,6 +335,7 @@ with gr.Blocks(css="style.css") as demo:
         use_upscaler,
         upscaler_strength,
         upscale_by,
+        json_input,  # Add JSON input to the list of inputs
     ]
 
     prompt.submit(
@@ -441,16 +345,11 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     ).then(
-        fn=generate,
+        fn=generate_and_update_history,  # Use the new function
         inputs=inputs,
-        outputs=[result, gr_metadata],
+        outputs=[result, gr_metadata, history_dropdown],  # Add history_dropdown to outputs
         api_name="run",
-    ).then(
-        fn=update_history,
-        inputs=[result, gr_metadata, history],
-        outputs=[history_dropdown, history],
     )
-
     negative_prompt.submit(
         fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
@@ -458,16 +357,11 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     ).then(
-        fn=generate,
+        fn=generate_and_update_history,  # Use the new function
         inputs=inputs,
-        outputs=[result, gr_metadata],
+        outputs=[result, gr_metadata, history_dropdown],  # Add history_dropdown to outputs
         api_name=False,
-    ).then(
-        fn=update_history,
-        inputs=[result, gr_metadata, history],
-        outputs=[history_dropdown, history],
     )
-
     run_button.click(
         fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
@@ -475,47 +369,25 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     ).then(
-        fn=generate,
+        fn=generate_and_update_history,  # Use the new function
         inputs=inputs,
-        outputs=[result, gr_metadata],
+        outputs=[result, gr_metadata, history_dropdown],  # Add history_dropdown to outputs
         api_name=False,
-    ).then(
-        fn=update_history,
-        inputs=[result, gr_metadata, history],
-        outputs=[history_dropdown, history],
-    )
-
-    apply_json_button.click(
-        fn=apply_json_parameters,
-        inputs=json_input,
-        outputs=[prompt, negative_prompt, seed, custom_width, custom_height,
-                 guidance_scale, num_inference_steps, sampler,
-                 aspect_ratio_selector, use_upscaler, upscaler_strength, upscale_by]
-    )
-
-    clear_button.click(
-        fn=lambda: (gr.update(value=""), gr.update(value=""), gr.update(value=0),
-                    gr.update(value=1024), gr.update(value=1024),
-                    gr.update(value=7.0), gr.update(value=30),
-                    gr.update(value="DPM++ 2M SDE Karras"),
-                    gr.update(value="1024 x 1024"), gr.update(value=False),
-                    gr.update(value=0.55), gr.update(value=1.5)),
-        inputs=[],
-        outputs=[prompt, negative_prompt, seed, custom_width, custom_height,
-                 guidance_scale, num_inference_steps, sampler,
-                 aspect_ratio_selector, use_upscaler, upscaler_strength, upscale_by]
     )
 
-    random_prompt_button.click(
-        fn=get_random_prompt,
-        inputs=[],
-        outputs=prompt
+    # Add event handler for generate_from_json button
+    generate_from_json.click(
+        fn=generate_and_update_history,
+        inputs=inputs,
+        outputs=[result, gr_metadata, history_dropdown],
+        api_name=False,
     )
 
+    # Add event handler for history_dropdown
     history_dropdown.change(
-        fn=lambda x, history: next((h["prompt"] for h in history if h["prompt"] == x), ""),
-        inputs=[history_dropdown, history],
-        outputs=prompt
+        fn=display_history_item,
+        inputs=[history_dropdown],
+        outputs=[history_image, history_metadata],
     )
 
-demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)
+demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)