refactored app.py

app.py

@@ -3,9 +3,7 @@ import torch
 import random
 import time
 from PIL import Image
-from utils import StableDiffusionConfig, StableDiffusionModels, generate_image
-
-# Set page config
+from custom_stable_diffusion import StableDiffusionConfig, StableDiffusionModels,ImageProcessor, generate_with_multiple_concepts,generate_with_concept_and_color,
 st.set_page_config(
     page_title="Butterfly Color Diffusion",
     page_icon="🦋",
@@ -32,10 +30,11 @@ def load_models():
         max_length=77
     )
     models = StableDiffusionModels(config)
+    image_processor = ImageProcessor(models, config)
     with st.spinner("Loading Stable Diffusion models... This may take a minute."):
         models.load_models()
         models.set_timesteps()
-    return models, config
+    return models, config, image_processor
 
 # Title and description
 st.title("🦋 Butterfly Color Diffusion")
@@ -55,6 +54,32 @@ prompt = st.sidebar.text_area(
     height=100
 )
 
+# Add concept selection dropdown
+available_concepts = [
+    "None (No concept)",
+    "concept-art-2-1",
+    "canna-lily-flowers102",
+    "arcane-style-jv",
+    "seismic-image",
+    "azalea-flowers102",
+    "photographic",
+    "realistic",
+    "detailed",
+    "national-geographic",
+    "macro-photography",
+    "nature-photography"
+]
+
+selected_concept = st.sidebar.selectbox(
+    "Select Concept Style",
+    available_concepts,
+    index=0,
+    help="Choose a concept style to apply to your image. Select 'None' to use standard generation."
+)
+
+# Convert "None" selection to actual None value
+concept_name = None if selected_concept == "None (No concept)" else selected_concept
+
 steps = st.sidebar.slider("Inference Steps", min_value=10, max_value=100, value=30, step=1)
 guidance_scale = st.sidebar.slider("Guidance Scale", min_value=1.0, max_value=15.0, value=7.5, step=0.1)
 seed = st.sidebar.number_input("Seed (0 for random)", min_value=0, max_value=1000000, value=0, step=1)
@@ -79,7 +104,7 @@ with col2:
 # Load models when needed
 if standard_button or color_button:
     if st.session_state.models is None:
-        st.session_state.models, st.session_state.config = load_models()
+        st.session_state.models, st.session_state.config ,st.session_state.image_processor= load_models()
     
     # Update config with current settings
     st.session_state.config.num_inference_steps = steps
@@ -98,17 +123,23 @@ if standard_button:
             progress_bar = st.progress(0)
             start_time = time.time()
             
-            image = generate_image(
+            image = generate_with_multiple_concepts(
                 models=st.session_state.models,
                 config=st.session_state.config,
+                image_processor=st.session_state.image_processor,
                 prompt=prompt,
+                concept_name=concept_name,  # Pass the selected concept
                 blue_loss_scale=0,
                 yellow_loss_scale=0,
+                guidance_interval=guidance_interval,
                 progress_bar=progress_bar
             )
             
             end_time = time.time()
-            st.image(image, caption="Standard Stable Diffusion", use_column_width=True)
+            caption = f"Standard Stable Diffusion"
+            if concept_name:
+                caption += f" with {concept_name} concept"
+            st.image(image, caption=caption, use_column_width=True)
             st.write(f"Generation time: {end_time - start_time:.2f} seconds")
 
 # Generate color-guided image
@@ -122,6 +153,7 @@ if color_button:
                 models=st.session_state.models,
                 config=st.session_state.config,
                 prompt=prompt,
+                concept_name=concept_name,  # Pass the selected concept
                 blue_loss_scale=0,
                 yellow_loss_scale=yellow_strength,
                 guidance_interval=guidance_interval,
@@ -129,7 +161,10 @@ if color_button:
             )
             
             end_time = time.time()
-            st.image(image, caption="Color-Guided Stable Diffusion", use_column_width=True)
+            caption = f"Color-Guided Stable Diffusion"
+            if concept_name:
+                caption += f" with {concept_name} concept"
+            st.image(image, caption=caption, use_column_width=True)
             st.write(f"Generation time: {end_time - start_time:.2f} seconds")
 
 # Explanation section
@@ -146,6 +181,9 @@ The color-guided approach adds a custom loss function during the diffusion proce
 
 This combination creates a yellow tone in the final image. The strength parameter controls how strongly this color guidance affects the generation process.
 
+### Concept Styles
+The concept styles use textual inversion embeddings to guide the image generation toward a particular artistic style or subject matter. These concepts have been trained on specific images and can dramatically change the look of your generated images.
+
 ### Technical Details
 During each step of the diffusion process, we:
 1. Calculate the predicted image at that step