app.py

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
import sys
import transformers
import diffusers
st.set_page_config(
    page_title="Butterfly Color Diffusion",
    page_icon="🦋",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state for models
if 'models' not in st.session_state:
    st.session_state.models = None
    st.session_state.config = None

# Add this near the top of your app.py
debug_mode = st.sidebar.checkbox("Debug Mode", value=True)

# Function to load models
@st.cache_resource
def load_models():
    if debug_mode:
        st.write("Debug: Starting model loading")
        st.write(f"Debug: Python version: {sys.version}")
        st.write(f"Debug: Torch version: {torch.__version__}")
        st.write(f"Debug: Transformers version: {transformers.__version__}")
        st.write(f"Debug: Diffusers version: {diffusers.__version__}")
    
    config = StableDiffusionConfig(
        height=512,
        width=512,
        num_inference_steps=30,
        guidance_scale=7.5,
        seed=random.randint(1, 1000000),
        batch_size=1,
        device=None,
        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()
    
    if debug_mode:
        st.write(f"Debug: Models loaded successfully. Device: {config.device}")
    
    return models, config, image_processor

# Title and description
st.title("🦋 Butterfly Color Diffusion")
st.markdown("""
Generate beautiful butterfly images with Stable Diffusion and explore color guidance technology 
that enhances yellow tones. Compare standard image generation with color-guided generation to see 
how targeted color loss functions can transform your results.
""")

# Sidebar with controls
st.sidebar.title("Generation Settings")

# Common settings
prompt = st.sidebar.text_area(
    "Prompt", 
    value="A detailed photograph of a colorful monarch butterfly with orange and black wings, resting on a purple flower in a lush garden with sunlight",
    height=100
)

# Add concept selection dropdown
available_concepts = [    
    "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)

# Color guidance settings
st.sidebar.markdown("---")
st.sidebar.subheader("Color Guidance Settings")
yellow_strength = st.sidebar.slider("Yellow Strength", min_value=0, max_value=500, value=200, step=10)
guidance_interval = st.sidebar.slider("Guidance Interval", min_value=1, max_value=10, value=5, step=1)

# Create two columns for the images
col1, col2 = st.columns(2)

with col1:
    st.subheader("Standard Stable Diffusion")
    standard_button = st.button("Generate Standard Image", use_container_width=True)

with col2:
    st.subheader("Color-Guided Stable Diffusion")
    color_button = st.button("Generate Color-Guided Image", use_container_width=True)

# 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 ,st.session_state.image_processor= load_models()
    
    # Update config with current settings
    st.session_state.config.num_inference_steps = steps
    st.session_state.config.guidance_scale = guidance_scale
    
    # Set seed
    if seed == 0:
        seed = random.randint(1, 1000000)
    st.session_state.config.seed = seed
    st.sidebar.write(f"Using seed: {seed}")

# Generate standard image
if standard_button:
    with col1:
        with st.spinner("Generating standard image..."):
            progress_bar = st.progress(0)
            start_time = time.time()
            
            # 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 [],
                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}")
                        st.stop()  # Stop execution here
                except Exception as e:
                    st.error(f"Error loading image: {str(e)}")
                    st.stop()  # Stop execution here
            
            caption = f"Standard Stable Diffusion"
            if concept_name:
                caption += f" with {concept_name} concept"
            st.image(image, caption=caption, use_container_width=True)
            st.write(f"Generation time: {end_time - start_time:.2f} seconds")

# Generate color-guided image
if color_button:
    with col2:
        with st.spinner("Generating color-guided image..."):
            progress_bar = st.progress(0)
            start_time = time.time()
            
            # 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,
                prompt=prompt,
                concepts=[concept_name] if concept_name else [],
                output_dir="concept_images",
                blue_loss_scale=0,
                yellow_loss_scale=yellow_strength                
            )
            
            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}")
                        st.stop()  # Stop execution here
                except Exception as e:
                    st.error(f"Error loading image: {str(e)}")
                    st.stop()  # Stop execution here
            
            caption = f"Color-Guided Stable Diffusion"
            if concept_name:
                caption += f" with {concept_name} concept"
            st.image(image, caption=caption, use_container_width=True)
            st.write(f"Generation time: {end_time - start_time:.2f} seconds")

# Explanation section
st.markdown("---")
st.header("How It Works")
st.markdown("""
### Standard Stable Diffusion
The standard approach uses text-to-image generation with classifier-free guidance to create images based on your prompt.

### Color-Guided Stable Diffusion
The color-guided approach adds a custom loss function during the diffusion process that encourages:
- Higher values in the red and green channels
- Lower values in the blue channel

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
2. Measure how far it is from our desired color profile
3. Calculate the gradient of this loss with respect to the latents
4. Adjust the latents to reduce the loss
5. Continue with the standard diffusion process

This approach allows for targeted control of specific visual attributes while maintaining the overall quality and coherence of the generated image.
""")

# Footer
st.markdown("---")
st.markdown("Created with ❤️ using Stable Diffusion and Streamlit")