refactored app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import torch
 import random
 import time
+import os
 from PIL import Image
 from custom_stable_diffusion import StableDiffusionConfig, StableDiffusionModels,ImageProcessor, generate_with_multiple_concepts,generate_with_multiple_concepts_and_color
 st.set_page_config(
@@ -123,16 +124,39 @@ if standard_button:
             progress_bar = st.progress(0)
             start_time = time.time()
             
-            image = generate_with_multiple_concepts(
+            # Call the generation function
+            result = generate_with_multiple_concepts(
                 models=st.session_state.models,
                 config=st.session_state.config,
                 image_processor=st.session_state.image_processor,
                 prompt=prompt,
-                concepts=[concept_name] if concept_name else [],  # Pass the selected concept
+                concepts=[concept_name] if concept_name else [],
                 output_dir="concept_images"
             )
             
             end_time = time.time()
+            
+            # Check if we got a valid image back
+            if result is not None and hasattr(result, 'format'):
+                # It's a PIL Image object
+                image = result
+            else:
+                # Try to load the image from the expected output path
+                try:
+                    if concept_name:
+                        image_path = f"concept_images/{concept_name}/{concept_name}.png"
+                    else:
+                        image_path = "concept_images/standard_image.png"
+                    
+                    if os.path.exists(image_path):
+                        image = Image.open(image_path)
+                    else:
+                        st.error(f"Could not find generated image at {image_path}")
+                        continue
+                except Exception as e:
+                    st.error(f"Error loading image: {str(e)}")
+                    continue
+            
             caption = f"Standard Stable Diffusion"
             if concept_name:
                 caption += f" with {concept_name} concept"
@@ -146,7 +170,8 @@ if color_button:
             progress_bar = st.progress(0)
             start_time = time.time()
             
-            image = generate_with_multiple_concepts_and_color(
+            # Call the generation function
+            result = generate_with_multiple_concepts_and_color(
                 models=st.session_state.models,
                 config=st.session_state.config,
                 image_processor=st.session_state.image_processor,
@@ -158,6 +183,30 @@ if color_button:
             )
             
             end_time = time.time()
+            
+            # Check if we got a valid image back
+            if result is not None and hasattr(result, 'format'):
+                # It's a PIL Image object
+                image = result
+            else:
+                # Try to load the image from the expected output path
+                try:
+                    if concept_name:
+                        # Determine the filename based on color guidance
+                        color_info = f"_yellow{yellow_strength}" if yellow_strength > 0 else ""
+                        image_path = f"concept_images/{concept_name}/{concept_name}{color_info}.png"
+                    else:
+                        image_path = "concept_images/color_guided_image.png"
+                    
+                    if os.path.exists(image_path):
+                        image = Image.open(image_path)
+                    else:
+                        st.error(f"Could not find generated image at {image_path}")
+                        continue
+                except Exception as e:
+                    st.error(f"Error loading image: {str(e)}")
+                    continue
+            
             caption = f"Color-Guided Stable Diffusion"
             if concept_name:
                 caption += f" with {concept_name} concept"