Update app.py

app.py

@@ -5,6 +5,9 @@ import numpy as np
 from transformers import CLIPTextModel, CLIPTokenizer
 from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
 from tqdm import tqdm
+import os
+import json
+import glob
 
 # Set device and dtype
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -20,7 +23,41 @@ text_encoder = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder
 unet = UNet2DConditionModel.from_pretrained(model_name, subfolder="unet", torch_dtype=dtype).to(device)
 scheduler = PNDMScheduler.from_pretrained(model_name, subfolder="scheduler")
 
-def generate_image(prompt, height, width, num_inference_steps, guidance_scale, seed):
+def load_examples_from_directory(sample_output_dir="sample_output"):
+    """
+    Load example data from the sample_output directory.
+    Assumes each image has a corresponding .json file with metadata.
+    """
+    examples = []
+    # Look for .json files in the directory
+    json_files = glob.glob(os.path.join(sample_output_dir, "*.json"))
+    
+    for json_file in json_files:
+        try:
+            with open(json_file, 'r') as f:
+                metadata = json.load(f)
+            # Ensure required fields are present
+            required_keys = ["prompt", "height", "width", "num_inference_steps", "guidance_scale", "seed"]
+            if all(key in metadata for key in required_keys):
+                examples.append([
+                    metadata["prompt"],
+                    metadata["height"],
+                    metadata["width"],
+                    metadata["num_inference_steps"],
+                    metadata["guidance_scale"],
+                    metadata["seed"]
+                ])
+        except Exception as e:
+            print(f"Error loading {json_file}: {e}")
+    
+    # If no valid examples are found, return a default example
+    if not examples:
+        examples = [
+            ["a serene landscape in Ghibli style", 64, 64, 50, 3.5, 42]
+        ]
+    return examples
+
+def generate_image(prompt, height, width, num_inference_steps, guidance_scale, seed, random_seed):
     # Validate inputs
     # if not prompt:
     #     return None, "Prompt cannot be empty."
@@ -38,7 +75,9 @@ def generate_image(prompt, height, width, num_inference_steps, guidance_scale, s
     # Set batch size
     batch_size = 1
 
-    # Create generator
+    # Handle random seed
+    if random_seed:
+        seed = torch.randint(0, 4294967295, (1,)).item()
     generator = torch.Generator(device=device).manual_seed(int(seed))
 
     # Tokenize and encode prompt
@@ -96,13 +135,14 @@ def generate_image(prompt, height, width, num_inference_steps, guidance_scale, s
     image = (image * 255).round().astype("uint8")
     pil_image = Image.fromarray(image[0])
 
-    return pil_image, "Image generated successfully!"
+    return pil_image, f"Image generated successfully! Seed used: {seed}"
 
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Ghibli-Style Image Generator")
-    gr.Markdown("Generate images in Ghibli style using a fine-tuned Stable Diffusion model. Enter a prompt and adjust parameters to create your image.")
-
+    gr.Markdown("Generate images in Ghibli style using a fine-tuned Stable Diffusion model. Enter ABOVE a prompt and adjust parameters to create your image.")
+    gr.Markdown("**Note:** For CPU inference, execution time is long (e.g., for 64x64 resolution with 50 inference steps, time is approximately 1800 seconds).")
+    
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label="Prompt", placeholder="e.g., 'a serene landscape in Ghibli style'")
@@ -111,15 +151,26 @@ with gr.Blocks() as demo:
             num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=50)
             guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=3.5)
             seed = gr.Slider(label="Seed", minimum=0, maximum=4294967295, step=1, value=42)
+            random_seed = gr.Checkbox(label="Use Random Seed", value=False)
             generate_btn = gr.Button("Generate Image")
         with gr.Column():
             output_image = gr.Image(label="Generated Image")
             output_text = gr.Textbox(label="Status")
 
+    gr.Markdown("### Example Prompts")
+    # Load examples from sample_output directory
+    examples_data = load_examples_from_directory("sample_output")
+    examples = gr.Dataframe(
+        value=examples_data,
+        headers=["Prompt", "Height", "Width", "Inference Steps", "Guidance Scale", "Seed"],
+        label="Examples"
+    )
+
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, guidance_scale, seed],
+        inputs=[prompt, height, width, num_inference_steps, guidance_scale, seed, random_seed],
         outputs=[output_image, output_text]
     )
 
-demo.launch()
+# Launch with limited users
+demo.launch(max_threads=3)