Update app.py

app.py

@@ -23,10 +23,11 @@ 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 load_examples_from_directory(sample_output_dir="sample_output"):
+def load_examples_from_directory(sample_output_dir="assets/example"):
     """
-    Load example data from the sample_output directory.
+    Load example data from the assets/example directory.
     Assumes each image has a corresponding .json file with metadata.
+    Returns a list of [prompt, height, width, num_inference_steps, guidance_scale, seed, json_filename].
     """
     examples = []
     json_files = glob.glob(os.path.join(sample_output_dir, "*.json"))
@@ -43,14 +44,15 @@ def load_examples_from_directory(sample_output_dir="sample_output"):
                     metadata["width"],
                     metadata["num_inference_steps"],
                     metadata["guidance_scale"],
-                    metadata["seed"]
+                    metadata["seed"],
+                    os.path.basename(json_file)  # Store the JSON filename for image lookup
                 ])
         except Exception as e:
             print(f"Error loading {json_file}: {e}")
     
     if not examples:
         examples = [
-            ["a serene landscape in Ghibli style", 64, 64, 50, 3.5, 42]
+            ["a serene landscape in Ghibli style", 64, 64, 50, 3.5, 42, None]
         ]
     return examples
 
@@ -122,25 +124,60 @@ def generate_image(prompt, height, width, num_inference_steps, guidance_scale, s
 
     return pil_image, f"Image generated successfully! Seed used: {seed}"
 
+def load_existing_image(json_filename, sample_output_dir="assets/example"):
+    """
+    Load an existing image corresponding to the given JSON filename.
+    Assumes the image has the same base name as the JSON file (e.g., example1.json -> example1.png).
+    """
+    if not json_filename:
+        return None, "No associated image found."
+    
+    # Get the base name without extension
+    base_name = os.path.splitext(json_filename)[0]
+    # Look for common image extensions
+    image_extensions = ["*.png", "*.jpg", "*.jpeg"]
+    image_path = None
+    
+    for ext in image_extensions:
+        candidates = glob.glob(os.path.join(sample_output_dir, f"{base_name}{ext}"))
+        if candidates:
+            image_path = candidates[0]
+            break
+    
+    if not image_path:
+        return None, f"No image found for {json_filename} in {sample_output_dir}."
+    
+    try:
+        pil_image = Image.open(image_path)
+        return pil_image, f"Loaded existing image for {json_filename}."
+    except Exception as e:
+        return None, f"Error loading image {image_path}: {e}"
+
 def on_example_select(selected_row):
     """
     Handle selection of a row in the examples Dataframe.
-    Returns the parameters to populate the input fields and trigger image generation.
+    Loads the existing image and updates input fields.
     """
+    prompt = selected_row[0]
+    height = selected_row[1]
+    width = selected_row[2]
+    num_inference_steps = selected_row[3]
+    guidance_scale = selected_row[4]
+    seed = selected_row[5]
+    json_filename = selected_row[6]  # JSON filename for image lookup
+    random_seed = False
+
+    # Load the existing image
+    image, status = load_existing_image(json_filename)
+    
     return (
-        selected_row[0],  # prompt
-        selected_row[1],  # height
-        selected_row[2],  # width
-        selected_row[3],  # num_inference_steps
-        selected_row[4],  # guidance_scale
-        selected_row[5],  # seed
-        False            # random_seed
+        prompt, height, width, num_inference_steps, guidance_scale, seed, random_seed, image, status
     )
 
 # 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 or select an example below to create your image.")
+    gr.Markdown("Generate images in Ghibli style using a fine-tuned Stable Diffusion model. Enter a prompt or select an example below to load a pre-generated 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():
@@ -158,10 +195,10 @@ with gr.Blocks() as demo:
             output_text = gr.Textbox(label="Status")
 
     gr.Markdown("### Example Prompts")
-    examples_data = load_examples_from_directory("sample_output")
+    examples_data = load_examples_from_directory("assets/example")
     examples = gr.Dataframe(
         value=examples_data,
-        headers=["Prompt", "Height", "Width", "Inference Steps", "Guidance Scale", "Seed"],
+        headers=["Prompt", "Height", "Width", "Inference Steps", "Guidance Scale", "Seed", "JSON File"],
         label="Examples",
         interactive=True
     )
@@ -170,11 +207,7 @@ with gr.Blocks() as demo:
     examples.select(
         fn=on_example_select,
         inputs=examples,
-        outputs=[prompt, height, width, num_inference_steps, guidance_scale, seed, random_seed]
-    ).then(
-        fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, guidance_scale, seed, random_seed],
-        outputs=[output_image, output_text]
+        outputs=[prompt, height, width, num_inference_steps, guidance_scale, seed, random_seed, output_image, output_text]
     )
 
     generate_btn.click(